Compare commits

..
Author SHA1 Message Date
lishunyang12 18252b9d8d [Lance] Fix x2t cache-offset SDPA mask for Q != K
PyTorch SDPA's is_causal=True with Q != K uses tril(diagonal=0) which is
top-left aligned. For a single decode step (Q=1) against a long ViT+text
cache (K=574+), only q_0 -> k_0 is unmasked — producing degenerate output
("TheTheThe..." token loops).

Bypass DiffusionAttention (whose _maybe_reshape_attn_mask only handles
2-D (B, K) masks) and call torch.nn.functional.scaled_dot_product_attention
directly with an explicit 4-D (1, 1, Q, K) bottom-right-aligned causal mask
when cache_len > 0. Also cast q/k to fp32 before RMSNorm for precision.

Verified all 7 tasks on Lance:
- x2t_image:  "No" (valid VQA answer)
- x2t_video:  "The purple sphere ... right direc..." (valid VQA answer)
- t2i / i2v / t2v / image_edit / video_edit all produce expected media.

Signed-off-by: lishunyang12 <lishunyang12@163.com>
2026-05-31 14:59:49 +00:00
lishunyang12 a326cd0d57 [Lance] Restore deploy YAML + patch post-main-merge API drift
After merging main into ``lance-model-support`` (commit c15f51a2),
Bagel's internal API was aggressively trimmed: ``forward_cache_update_text``,
``forward_cache_update_vit``, ``generate_image``, ``generate_text``, and
``Qwen2MoTForCausalLM.forward`` no longer accept ``packed_text_indexes`` /
``packed_key_value_indexes`` / ``key_values_lens`` / ``packed_query_indexes``
— they're now derived from ``past_key_values``.  Lance's pipeline was
still passing those kwargs and crashed on every task.

Changes:

- ``vllm_omni/deploy/lance.yaml`` (restored): Lance's HF ``config.json``
  is descriptive metadata with no ``model_type`` field, so the loader
  can't auto-detect the pipeline.  This YAML is the minimal hook that
  tells the loader "use the Lance pipeline" — one stage, no multi-stage
  orchestration.
- ``examples/offline_inference/lance/gradio_demo.py``: pass
  ``deploy_config`` again so ``Omni`` picks up the YAML.
- ``vllm_omni/diffusion/models/bagel/bagel_transformer.py``: in the
  batched CFG setup, concat ``branches_pid`` along ``dim=-1`` (the
  sequence axis) instead of the default ``dim=0`` so Lance's 3-D mRoPE
  position ids ``(3, S)`` aren't collapsed into ``(6, S)`` —
  ``BagelRotaryEmbedding`` enforces exactly 3 axes for mRoPE.
- ``vllm_omni/diffusion/models/bagel/pipeline_bagel.py``: mirror the
  generated PIL image into ``custom_output["image"]`` so callers reading
  through the orchestrator IPC boundary (which strips the bare
  ``output`` field) can still recover it.  ``video_frames`` already used
  this pattern.
- ``vllm_omni/diffusion/models/lance/lance_transformer.py``: drop
  ``packed_query_indexes`` / ``key_values_lens`` /
  ``packed_key_value_indexes`` from the ``language_model.forward`` call
  inside ``forward_cache_update_vae`` — derived now.
- ``vllm_omni/diffusion/models/lance/pipeline_lance.py``:
    * drop the same 3 fields from the inline ``inp`` dicts that build
      ``forward_cache_update_text`` arguments in ``_forward_image_edit``,
      ``_forward_i2v``, ``_forward_video_edit``, and
      ``_x2t_raw_text_prefill``.
    * drop the 6 ``cfg_text_packed_query_indexes`` /
      ``cfg_text_key_values_lens`` / ``cfg_text_packed_key_value_indexes``
      (+ ``cfg_img_*`` triple) kwargs from the four ``generate_image``
      call sites.
    * filter ``key_values_lens`` / ``packed_indexes`` /
      ``packed_key_value_indexes`` out of the ``gen_input_lat`` /
      ``gen_input_vit`` dicts before unpacking into ``generate_image``
      and ``forward_cache_update_vit`` (the ``prepare_*_latent`` /
      ``prepare_vit_*`` helpers still return them for the
      backwards-compatible parent-class path).
    * drop ``key_values_lens`` / ``packed_key_value_indexes`` from
      ``_x2t_prepare_assistant_start``'s returned dict so
      ``Bagel.generate_text`` doesn't reject them.

Verification (local online inference, single GPU pair, CFG=4, 10 steps):

| Task          | Status                                          |
|---------------|-------------------------------------------------|
| t2i           | OK  (~2 s)                                      |
| image_edit    | OK  (~2.5 s)                                    |
| i2v           | OK  (~8 s @ 25 frames)                          |
| t2v           | OK  (~7 s @ 25 frames)                          |
| video_edit    | OK  (~60 s)                                     |
| x2t_image     | RUNS but generates degenerate token repetition  |
| x2t_video     | RUNS but generates degenerate token repetition  |

x2t degeneration is a deeper semantic mismatch with main's new
``past_key_values``-derived cache layout: Lance's x2t prefill stitches
the chat template across multiple ``forward_cache_update_*`` calls
(system → ViT → question → ``assistant\\n``), but the new API expects a
single contiguous KV layout.  Fixing requires reworking the segmented
prefill to match main's new cache contract — left as a follow-up since
the 5 generative tasks now work.

Signed-off-by: lishunyang12 <lishunyang12@163.com>
2026-05-31 13:28:11 +00:00
lishunyang12 c15f51a2f6 Merge branch 'main' into lance-model-support; drop lance.yaml
Resolves conflicts in:
- vllm_omni/diffusion/registry.py: keep both LancePipeline (HEAD)
  and MingImagePipeline (main) entries.
- vllm_omni/diffusion/models/bagel/bagel_transformer.py: take main's
  simpler cfg_batched payload; HEAD referenced undefined locals.

Also strips vllm_omni/deploy/lance.yaml -- Lance is single-stage
diffusion, so its engine knobs (pipeline=lance, enforce_eager,
trust_remote_code, max_num_seqs=1, ...) flow directly to Omni(...)
/ vllm-omni serve instead of through a deploy YAML.
create_default_diffusion builds the stage config from those kwargs.

Touched: examples/offline_inference/lance/{end2end,gradio_demo}.py,
tests/e2e/online_serving/test_lance.py, tests/helpers/stage_config.py.

Signed-off-by: lishunyang12 <lishunyang12@163.com>
2026-05-31 11:50:21 +00:00
lishunyang12 7c934b6188 [Lance] Add Gradio demo and align inference with upstream
- Add ``examples/offline_inference/lance/gradio_demo.py`` — a unified 7-task UI
  (t2i / t2v / i2v / image_edit / video_edit / x2t_image / x2t_video) that
  loads ``Lance_3B`` and ``Lance_3B_Video`` on two GPUs so image and video
  tasks run concurrently.  Surfaces official Lance examples per task; aspect
  ratio is auto-snapped to one of upstream's 6 buckets from the input
  image/video.
- Fix x2t_image / x2t_video output format: route x2t_image through a new
  ``_forward_x2t_image`` handler (was falling through to the img2img path
  and emitting garbage), build the chat prefill in three segments
  (prefix → ViT → suffix) so the ``<|im_end|>`` markers line up with
  upstream's template, and start generation from a newline token instead of
  ``<|im_start|>`` so the assistant turn doesn't close itself with an
  immediate EOS.  Add a ``system_prompt`` override to ``render_lance_prompt``
  so per-example QA instructions surface ("Look at the image carefully and
  answer the question.") instead of the caption-style default.
- ``_forward_image_edit``: also stash the PIL image into ``custom_output``
  so callers reading ``outputs[0].output`` *or* ``custom_output["image"]``
  both work — the orchestrator IPC strips the bare ``output`` field across
  process boundaries.
- ``pipeline_registry.py``, ``model_executor/models/lance/*``, registry
  cleanup: remove the dead Lance-specific autoregressive shim
  (``lance.py``) and its stage-input processor; the diffusion pipeline owns
  the full flow now.
- Aspect-ratio buckets aligned with upstream's ``BucketResize`` (6 ratios:
  21:9 / 16:9 / 4:3 / 1:1 / 3:4 / 9:16, sized for ``resolution_vae`` ∈
  {480, 640, 768}).  3:2 / 2:3 dropped — upstream snaps those to 4:3 / 3:4.
- Gradio defaults match upstream's official scripts (cfg_text_scale=4.0,
  validation_num_timesteps=30, validation_timestep_shift=3.5, seed=42); only
  the demo *duration* is intentionally shorter than upstream's 121-frame
  default to keep the UI snappy (user can extend to 10s via slider for
  byte-identical reproduction).

Signed-off-by: lishunyang12 <lishunyang12@163.com>

Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>
2026-05-31 11:40:40 +00:00
Hacoandzzhuoxin1508 74050cdacd [BugFix] Fix diffusion parallel_config YAML override and add deploy config field allowlist (#3483)
Signed-off-by: xiaohajiayou <923390377@qq.com>
Co-authored-by: zzhuoxin1508 <zzhuoxin1508@users.noreply.github.com>

Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>
2026-05-31 08:58:52 +00:00
SYLAR f0cec3ecd4 Merge branch 'main' into lance-model-support 2026-05-22 04:13:58 +08:00
lishunyang12 7b885f9d13 [Lance] Precision matching, recipe, and README cleanup
- Precision matching across the Lance transformer and pipeline so vLLM-Omni
  outputs match the upstream reference end-to-end (text2img / img2img /
  text2video / video2video / image and video understanding).
- Add recipes/ByteDance/Lance.md covering all single-stage modalities for
  offline + online; register it in recipes/README.md.
- Offline README: document --height/--width and update the default
  resolution to 1024x1024 (matches the upstream inference_lance.sh).
- Gitignore lance_compare/ (local precision-matching dump scripts; not part
  of the repo).
- Remove LANCE_PR.md; the PR description lives on GitHub.

Signed-off-by: lishunyang12 <lishunyang12@163.com>
2026-05-21 20:07:40 +00:00
lishunyang 69bcc6416d [New Model] Lance (ByteDance) — unified AR+diffusion multimodal (#3697)
Adds Lance (bytedance-research/Lance), a 3B unified AR+diffusion
multimodal model on a Qwen2.5-VL backbone. Lance is BAGEL-lineage:
it reuses vLLM-Omni's BAGEL Mixture-of-Transformers core verbatim and
specializes only the ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the
checkpoint layout.

All six single-stage Lance modalities verified end-to-end on 1x B300:
t2i, t2v, image_edit, video_edit, x2t_image (image understanding) and
x2t_video (video understanding). 3-D mRoPE is wired through
BagelRotaryEmbedding for the x2t / video / edit paths. Includes the
upstream gap-analysis quality fixes (task-specific system prompts, full
Qwen chat template, upstream inference defaults and aspect-ratio
buckets).

Ships an offline example (examples/offline_inference/lance/), an
OpenAI-compatible online-serving example
(examples/online_serving/lance/) and an e2e test
(tests/e2e/online_serving/test_lance.py).

See LANCE_PR.md for the architecture rationale, per-modality status and
the two remaining follow-ups (two-stage AR+DiT, video_edit polish).

Closes #3697.

Signed-off-by: lishunyang <lishunyang03@gmail.com>
2026-05-18 22:09:14 +00:00
517 changed files with 6982 additions and 64839 deletions

No files matched your search

+1 -1
View File
@@ -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.21.0"
priority: 100
timeout_in_minutes: 60
soft_fail: false
+4 -4
View File
@@ -1,14 +1,14 @@
# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
# exception in upload_pipeline.py (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
# Document 2: appended after `---`; same file, read by upload_pipeline.py (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
steps:
- label: ":github: Resolve skip-ci & upload pipeline"
key: upload-ci-pipeline
commands:
- "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml"
- "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
agents:
queue: "cpu_queue_premerge"
@@ -41,7 +41,7 @@ steps:
key: upload-merge-pipeline
if: __UPLOAD_MERGE_IF__
commands:
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
- buildkite-agent pipeline upload .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
+1 -1
View File
@@ -23,7 +23,7 @@ steps:
depends_on: image-build
key: upload-merge-pipeline
commands:
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
- buildkite-agent pipeline upload .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
-308
View File
@@ -1,308 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Render and optionally upload Buildkite pipeline YAML with diff-aware logic.
Bootstrap mode (pipeline.yml with __IMAGE_BUILD_IF__ placeholders):
- Detect docs-only PR/main changes and substitute skip-ci ``if`` expressions.
Test pipeline mode (e.g. test-merge.yml):
- Drop steps or groups whose ``source_file_dependencies`` do not match changed files.
- ``source_file_dependencies`` may be set on a leaf step or on a group step.
- Always strip ``source_file_dependencies`` from uploaded YAML (Buildkite does not
recognize this uploader-only key).
Usage:
python3 upload_pipeline.py [--upload] <pipeline.yml>
# Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
python3 upload_pipeline.py --upload .buildkite/pipeline.yml
# Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
python3 upload_pipeline.py --upload .buildkite/test-merge.yml
"""
from __future__ import annotations
import argparse
import os
import subprocess
import sys
from pathlib import Path
from typing import Any
try:
import yaml
except ImportError:
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "-q", "pyyaml"],
)
import yaml
LOG = "upload_pipeline"
ROOT = Path(__file__).resolve().parent.parent.parent
DOC_SEP = "\n---\n"
BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
def _log(message: str) -> None:
print(f"{LOG}: {message}", file=sys.stderr)
def _git(*args: str) -> subprocess.CompletedProcess[str]:
return subprocess.run(
["git", *args],
check=False,
capture_output=True,
text=True,
)
def resolve_changed_files() -> list[str] | None:
"""Return changed file paths, or None when diff cannot be resolved."""
is_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false" and os.environ.get(
"BUILDKITE_PULL_REQUEST", ""
)
commit = os.environ.get("BUILDKITE_COMMIT", "")
if is_pr:
base_branch = os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH", "main")
base_ref = f"origin/{base_branch}"
if _git("rev-parse", "--verify", base_ref).returncode != 0:
_log(f"origin/{base_branch} not found locally; trying fetch")
_git("fetch", "--depth=200", "origin", base_branch)
if _git("rev-parse", "--verify", base_ref).returncode != 0:
if _git("rev-parse", "--verify", base_branch).returncode == 0:
base_ref = base_branch
else:
_log(f"cannot resolve PR base {base_branch}; using safe defaults")
return None
diff_range = f"{base_ref}...{commit}"
elif os.environ.get("BUILDKITE_BRANCH", "") == "main":
if _git("rev-parse", "--verify", f"{commit}^").returncode != 0:
_log("main commit has no parent; using safe defaults")
return None
diff_range = f"{commit}^..{commit}"
else:
_log("not PR/main build; using safe defaults")
return None
result = _git("diff", "--name-only", diff_range)
if result.returncode != 0:
_log(f"git diff failed ({diff_range}); using safe defaults")
return None
files = [line for line in result.stdout.splitlines() if line.strip()]
_log(f"{len(files)} changed file(s)")
return files
def is_docs_only_change(changed_files: list[str]) -> bool:
has_any = False
for file_path in changed_files:
if not file_path:
continue
has_any = True
if file_path.startswith("docs/"):
continue
if file_path.endswith(".md"):
continue
if file_path == "mkdocs.yml":
continue
return False
return has_any
def resolve_skip_ci(changed_files: list[str] | None) -> bool:
if changed_files is None:
_log("skip-ci=0 (could not resolve changed files)")
return False
if is_docs_only_change(changed_files):
_log("docs-only change detected; skip-ci=1")
return True
_log("non-doc changes detected; skip-ci=0")
return False
def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
if DOC_SEP in text:
_, continuation = text.split(DOC_SEP, 1)
else:
continuation = text
nightly_only = (
'(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
)
if skip_ci:
image_if = f"'{nightly_only}'"
ready_if = "'false'"
merge_if = "'false'"
else:
image_if = "'true'"
ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\''
merge_if = (
'\'(build.branch == "main" && build.env("NIGHTLY") != "1" '
'&& build.env("WEEKLY") != "1") || '
'(build.branch != "main" && build.pull_request.labels includes "merge-test")\''
)
return (
continuation.replace("__IMAGE_BUILD_IF__", image_if)
.replace("__UPLOAD_READY_IF__", ready_if)
.replace("__UPLOAD_MERGE_IF__", merge_if)
)
def _matches_dependencies(changed_files: list[str], prefixes: list[str]) -> bool:
for path in changed_files:
for prefix in prefixes:
normalized = prefix.rstrip("/")
if path == normalized or path.startswith(f"{normalized}/"):
return True
return False
def _strip_source_file_dependencies(step: dict[str, Any]) -> dict[str, Any]:
return {key: value for key, value in step.items() if key != "source_file_dependencies"}
def _step_label(step: dict[str, Any]) -> str:
return str(step.get("group") or step.get("label") or "<step>")
def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
filtered: list[Any] = []
for step in steps:
if not isinstance(step, dict):
filtered.append(step)
continue
deps = step.get("source_file_dependencies")
if deps is not None and not isinstance(deps, list):
raise ValueError(
f"source_file_dependencies must be a list in step {_step_label(step)!r}",
)
if deps is not None and not _matches_dependencies(changed_files, deps):
_log(
f"skip {_step_label(step)!r} (no changes under {deps})",
)
continue
nested = step.get("steps")
if nested is not None:
kept_nested = _filter_steps(nested, changed_files)
if not kept_nested:
_log(f"omit empty group {_step_label(step)!r}")
continue
new_step = _strip_source_file_dependencies(step)
new_step["steps"] = kept_nested
filtered.append(new_step)
continue
if deps is not None:
filtered.append(_strip_source_file_dependencies(step))
else:
filtered.append(step)
return filtered
def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
"""Remove uploader-only keys while keeping all steps (no diff filtering)."""
stripped: list[Any] = []
for step in steps:
if not isinstance(step, dict):
stripped.append(step)
continue
deps = step.get("source_file_dependencies")
if deps is not None and not isinstance(deps, list):
raise ValueError(
f"source_file_dependencies must be a list in step {_step_label(step)!r}",
)
nested = step.get("steps")
new_step = _strip_source_file_dependencies(step)
if nested is not None:
new_step["steps"] = _strip_uploader_metadata_from_steps(nested)
stripped.append(new_step)
return stripped
def render_test_pipeline(
doc: dict[str, Any],
changed_files: list[str] | None,
) -> dict[str, Any]:
steps = doc.get("steps")
if not isinstance(steps, list):
return doc
if changed_files is not None:
steps = _filter_steps(steps, changed_files)
else:
steps = _strip_uploader_metadata_from_steps(steps)
return {**doc, "steps": steps}
def resolve_pipeline_path(arg: str) -> Path:
path = Path(arg)
if path.is_absolute():
return path
return ROOT / path
def render_pipeline(path: Path) -> str:
text = path.read_text(encoding="utf-8")
changed_files = resolve_changed_files()
if BOOTSTRAP_MARKER in text:
rendered = render_bootstrap_pipeline(text, skip_ci=resolve_skip_ci(changed_files))
return rendered
doc = yaml.safe_load(text)
if not isinstance(doc, dict):
raise ValueError(f"invalid pipeline YAML: {path}")
doc = render_test_pipeline(doc, changed_files)
return yaml.safe_dump(doc, sort_keys=False)
def upload_to_buildkite(content: str) -> None:
subprocess.run(
["buildkite-agent", "pipeline", "upload"],
input=content,
text=True,
check=True,
)
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"pipeline",
nargs="?",
default=".buildkite/pipeline.yml",
help="Pipeline YAML path (default: .buildkite/pipeline.yml)",
)
parser.add_argument(
"--upload",
action="store_true",
help="Pipe rendered YAML to buildkite-agent pipeline upload",
)
args = parser.parse_args()
path = resolve_pipeline_path(args.pipeline)
if not path.is_file():
_log(f"missing pipeline file: {path}")
return 1
rendered = render_pipeline(path)
if args.upload:
upload_to_buildkite(rendered)
else:
sys.stdout.write(rendered)
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -0,0 +1,146 @@
#!/usr/bin/env bash
# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
#
# Usage: upload_pipeline_with_skip_ci.sh [pipeline_yaml]
# pipeline_yaml: path relative to repo root or absolute (default: .buildkite/pipeline.yml)
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
PIPELINE_ARG="${1:-.buildkite/pipeline.yml}"
if [[ "${PIPELINE_ARG}" = /* ]]; then
PIPELINE_YML="${PIPELINE_ARG}"
else
PIPELINE_YML="${ROOT}/${PIPELINE_ARG}"
fi
# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
is_docs_only_change() {
local file_path
local has_any=0
while IFS= read -r file_path; do
[[ -z "${file_path}" ]] && continue
has_any=1
if [[ "${file_path}" == docs/* ]]; then
continue
fi
if [[ "${file_path}" == *.md ]]; then
continue
fi
if [[ "${file_path}" == "mkdocs.yml" ]]; then
continue
fi
return 1
done
[[ "${has_any}" -eq 1 ]]
}
resolve_skip_ci() {
local is_pr_build=0
local files
local base_branch base_ref
if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
is_pr_build=1
fi
if [[ "${is_pr_build}" -eq 1 ]]; then
base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
fi
base_ref=""
if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
base_ref="origin/${base_branch}"
elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
base_ref="${base_branch}"
else
echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
echo -n 0
return 0
fi
if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
echo -n 0
return 0
fi
elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
echo -n 0
return 0
fi
if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
echo -n 0
return 0
fi
else
echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
echo -n 0
return 0
fi
if is_docs_only_change <<< "${files}"; then
echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
echo -n 1
return 0
fi
echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
echo -n 0
}
SKIP_CI="$(resolve_skip_ci)"
if [[ ! -f "${PIPELINE_YML}" ]]; then
echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
exit 1
fi
export ROOT SKIP_CI PIPELINE_YML
python3 <<'PY' | buildkite-agent pipeline upload
import os
import pathlib
path = pathlib.Path(os.environ["PIPELINE_YML"])
text = path.read_text(encoding="utf-8")
sep = "\n---\n"
# Two supported layouts:
# - multi-doc: doc 1 = bootstrap (loaded by Buildkite), doc 2 (after `---`) = uploaded steps.
# - single-doc: whole file is the uploaded steps; caller (UI init step) is the bootstrap.
if sep in text:
_, continuation = text.split(sep, 1)
else:
continuation = text
skip = os.environ.get("SKIP_CI") == "1"
# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
nightly_only = (
'(build.pull_request.labels includes "nightly-test") '
'|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
)
# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
if skip:
rep = f"'{nightly_only}'"
ready_rep = "'false'"
merge_rep = "'false'"
else:
rep = "'true'"
ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\" && build.env(\"WEEKLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
rendered = (
continuation
.replace("__IMAGE_BUILD_IF__", rep)
.replace("__UPLOAD_READY_IF__", ready_rep)
.replace("__UPLOAD_MERGE_IF__", merge_rep)
)
print(rendered, end="")
PY
+146 -160
View File
@@ -1,111 +1,41 @@
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HF_HUB_DOWNLOAD_TIMEOUT: 300
HF_HUB_ETAG_TIMEOUT: 60
steps:
- group: ":card_index_dividers: Simple Test"
steps:
- label: "Simple · Diffusion & Model Executor Test"
agent_pool: mi325_1
depends_on: amd-build
timeout_in_minutes: 50
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple Unit Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple · Other Test"
agent_pool: mi325_1
depends_on: amd-build
timeout_in_minutes: 50
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Diffusion Images API LoRA E2E"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
- group: ":card_index_dividers: Diffusion Test"
steps:
- label: "Diffusion · LoRA Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
agent_pool: mi325_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
# merge-only tests
- label: "Diffusion Tensor Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
- label: "Diffusion Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
# - label: "Audio Generation Model Test"
@@ -119,6 +49,52 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
- label: "Diffusion Cache Backend Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
agent_pool: mi325_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
# merge-only tests
- label: "Diffusion Tensor Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
- label: "Diffusion Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Engine Test"
agent_pool: mi325_1
depends_on: amd-build
@@ -132,65 +108,50 @@ steps:
# This model seems to be important, so I am trying to run it here.
# Since this is from test-nightly.yml,
# I keep this non-blocking.
- group: ":card_index_dividers: E2E Test"
steps:
- label: "Omni Model Test Qwen2-5-Omni"
agent_pool: mi325_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
- label: "Omni Model Test Qwen2-5-Omni"
agent_pool: mi325_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
- label: "Omni Model Test Qwen3-Omni"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
- label: "Omni Model Test Qwen3-Omni"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
- label: "Qwen3-TTS CustomVoice E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Qwen3-TTS CustomVoice E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Qwen3-TTS Base E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
- label: "Qwen3-TTS Base E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
'
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
@@ -231,3 +192,28 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - export VLLM_ROCM_USE_AITER_RMSNORM=0
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
+132 -155
View File
@@ -1,42 +1,15 @@
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HF_HUB_DOWNLOAD_TIMEOUT: 300
HF_HUB_ETAG_TIMEOUT: 60
steps:
- group: ":card_index_dividers: Simple Test"
steps:
- label: "Simple · Diffusion & Model Executor Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple · Other Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Custom Pipeline Test"
- label: "Simple Unit Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Voxtral TTS CUDA Unit Test"
agent_pool: mi325_1
@@ -46,83 +19,39 @@ steps:
commands:
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
- group: ":card_index_dividers: Diffusion Test"
steps:
- label: "Diffusion · Batch Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion Batching Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
- label: "Diffusion · Offloader Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- label: "Custom Pipeline Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
- label: "Diffusion · Cache Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
- label: "Diffusion Sequence Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
"
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
# - label: "Audio Generation Model Test"
@@ -136,6 +65,56 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
- label: "Diffusion Cache Backend Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
- label: "Diffusion Sequence Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
"
- label: "Engine Test"
agent_pool: mi325_1
depends_on: amd-build
@@ -159,56 +138,43 @@ steps:
# - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
# - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
- group: ":card_index_dividers: E2E Test"
steps:
- label: "MiMo-Audio E2E Test with H100"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
'
- label: "MiMo-Audio E2E Test with H100"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
'
- label: "Qwen3-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
- label: "Qwen3-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
- label: "CosyVoice3-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
'
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
@@ -248,3 +214,14 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - export VLLM_ROCM_USE_AITER_RMSNORM=0
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
- label: "CosyVoice3-TTS E2E Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
'
+5 -122
View File
@@ -44,9 +44,6 @@ steps:
depends_on: upload-merge-pipeline
steps:
- label: "Diffusion · LoRA Test"
source_file_dependencies:
- tests/diffusion/lora/
- vllm_omni/diffusion/lora/
timeout_in_minutes: 30
commands:
- pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
@@ -64,9 +61,6 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Distributed Test"
source_file_dependencies:
- tests/distributed/
- vllm_omni/distributed/
timeout_in_minutes: 60
commands:
- |
@@ -111,12 +105,9 @@ steps:
- group: ":card_index_dividers: Entrypoints Test"
depends_on: upload-merge-pipeline
source_file_dependencies:
- tests/entrypoints/
- vllm_omni/entrypoints/
steps:
- label: "Entrypoint Test with H100"
timeout_in_minutes: 60
timeout_in_minutes: 35
commands:
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
agents:
@@ -174,13 +165,6 @@ steps:
depends_on: upload-merge-pipeline
steps:
- label: "TTS · Qwen3-TTS CustomVoice Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
- tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
- vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml
commands:
- |
timeout 20m bash -c "
@@ -202,12 +186,6 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Qwen3-TTS Base Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_tts_base.py
- tests/e2e/offline_inference/test_qwen3_tts_base.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
commands:
- |
timeout 20m bash -c "
@@ -228,34 +206,7 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Ming-Omni-TTS Test"
timeout_in_minutes: 20
commands:
- |
timeout 20m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
pytest -s -v tests/e2e/online_serving/test_ming_tts.py tests/e2e/offline_inference/test_ming_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Voxtral-TTS Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_voxtral_tts.py
- tests/e2e/online_serving/test_voxtral_tts.py
- vllm_omni/model_executor/models/voxtral_tts/
- vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
- vllm_omni/deploy/voxtral_tts.yaml
timeout_in_minutes: 20
commands:
- |
@@ -298,12 +249,6 @@ steps:
type: DirectoryOrCreate
- label: "TTS · GLM-TTS Test"
source_file_dependencies:
- tests/e2e/online_serving/test_glm_tts.py
- tests/e2e/offline_inference/test_glm_tts.py
- vllm_omni/model_executor/models/glm_tts/
- vllm_omni/model_executor/stage_input_processors/glm_tts.py
- vllm_omni/deploy/glm_tts.yaml
timeout_in_minutes: 20
commands:
- |
@@ -346,12 +291,6 @@ steps:
type: DirectoryOrCreate
- label: "Omni · Qwen3-Omni Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_qwen3_omni.py
- tests/e2e/online_serving/test_qwen3_omni.py
- vllm_omni/model_executor/models/qwen3_omni/
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
- vllm_omni/deploy/qwen3_omni_moe.yaml
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -390,11 +329,6 @@ steps:
type: DirectoryOrCreate
- label: "Omni · Mimo Audio Test"
source_file_dependencies:
- tests/e2e/online_serving/test_mimo_audio.py
- vllm_omni/model_executor/models/mimo_audio/
- vllm_omni/model_executor/stage_input_processors/mimo_audio.py
- vllm_omni/deploy/mimo_audio.yaml
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -432,19 +366,14 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_voxcpm2_tts.py
- tests/e2e/online_serving/test_voxcpm2_tts.py
- vllm_omni/model_executor/models/voxcpm2/
- vllm_omni/deploy/voxcpm2.yaml
timeout_in_minutes: 30
- label: "Omni · VoxCPM2 Test"
timeout_in_minutes: 20
commands:
- |
timeout 30m bash -c "
timeout 20m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
agents:
queue: "gpu_1_queue"
@@ -461,12 +390,6 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni · Higgs-Audio V2 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_higgs_audio_v2.py
- tests/e2e/offline_inference/test_higgs_audio_v2.py
- vllm_omni/model_executor/models/higgs_audio_v2/
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py
- vllm_omni/deploy/higgs_audio_v2.yaml
timeout_in_minutes: 30
commands:
- |
@@ -492,9 +415,6 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · Qwen Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_edit.py
- vllm_omni/diffusion/models/qwen_image/
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -533,9 +453,6 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Longcat Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_longcat_image_edit.py
- vllm_omni/diffusion/models/longcat_image/
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/online_serving/test_longcat_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -574,9 +491,6 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Longcat Image Test"
source_file_dependencies:
- tests/e2e/online_serving/test_longcat_image.py
- vllm_omni/diffusion/models/longcat_image/
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/online_serving/test_longcat_image.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -615,13 +529,6 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Bagel Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_bagel.py
- tests/e2e/online_serving/test_bagel.py
- vllm_omni/model_executor/models/bagel/
- vllm_omni/diffusion/models/bagel/
- vllm_omni/model_executor/stage_input_processors/bagel.py
- vllm_omni/deploy/bagel.yaml
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -661,9 +568,6 @@ steps:
- label: "Diffusion · Qwen Image Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 40m bash -c '
@@ -705,9 +609,6 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Qwen Image Layered Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_layered.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 40m bash -c "
@@ -749,10 +650,6 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Wan22 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_wan22_t2v.py
- tests/e2e/online_serving/test_wan22_t2v.py
- vllm_omni/diffusion/models/wan2_2/
commands:
- |
timeout 40m bash -c "
@@ -794,9 +691,6 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · HunyuanVideo 1.5 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_hunyuan_video_15.py
- vllm_omni/diffusion/models/hunyuan_video/
timeout_in_minutes: 40
commands:
- |
@@ -838,10 +732,6 @@ steps:
type: DirectoryOrCreate
- label: "Diffusion · Z Image Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_z_image.py
- tests/e2e/online_serving/test_z_image.py
- vllm_omni/diffusion/models/z_image/
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
@@ -859,10 +749,6 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · FLUX.2 Klein Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_flux2_klein.py
- tests/e2e/online_serving/test_flux2_klein.py
- vllm_omni/diffusion/models/flux2_klein/
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
@@ -880,9 +766,6 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · SD3 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_sd3.py
- vllm_omni/diffusion/models/sd3/
timeout_in_minutes: 45
commands:
- |
+4 -189
View File
@@ -394,11 +394,11 @@ steps:
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
steps:
- label: ":full_moon: TTS · Function Test with L4"
- label: ":full_moon: TTS · Function Test"
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
agents:
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -413,69 +413,6 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: ":full_moon: TTS · Function Test with H100"
# timeout_in_minutes: 90
# commands:
# - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
# - pytest -s -v tests/e2e/ -m "full_model and H100 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
# 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
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_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: ":full_moon: TTS · Perf Test"
key: nightly-tts-performance
timeout_in_minutes: 180
@@ -522,8 +459,8 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
key: nightly-tts-performance-voxcpm2
- label: ":full_moon: VoxCPM2 · Perf Test"
key: nightly-voxcpm2-performance
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
@@ -591,7 +528,6 @@ steps:
tests/e2e/online_serving/test_longcat_image_edit_expansion.py
tests/e2e/online_serving/test_flux_2_dev_expansion.py
tests/e2e/online_serving/test_bagel_expansion.py
tests/e2e/offline_inference/test_sensenova_u1_expansion.py
-m "full_model and diffusion and H100 and not distributed_cuda"
--run-level "full_model"
agents:
@@ -759,87 +695,6 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU DreamZero"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_dreamzero_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
timeout_in_minutes: 120
commands:
- pytest -sv tests/e2e/online_serving/test_glm_image_expansion.py -m "full_model and diffusion and H100" --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
timeout_in_minutes: 60
@@ -938,46 +793,6 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
timeout_in_minutes: 180
commands:
- export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct"
- export HUNYUAN_IMAGE3_DEVICES="0,1,2,3"
- pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_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: 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: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
key: nightly-diffusion-x2iat-performance-qwen-image
timeout_in_minutes: 180
+5 -64
View File
@@ -183,7 +183,7 @@ steps:
depends_on: upload-ready-pipeline
commands:
- |
timeout 20m bash -c "
timeout 15m bash -c "
pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
"
agents:
@@ -273,50 +273,14 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni · Qwen3-Omni Realtime WebSocket"
- label: "Omni · VoxCPM2 Test"
timeout_in_minutes: 20
commands:
- |
timeout 30m bash -c "
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -m 'core_model' --run-level 'core_model'
"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "TTS · VoxCPM2 Test"
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
timeout 20m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'core_model' --run-level 'core_model'
"
agents:
queue: "gpu_1_queue"
@@ -376,29 +340,6 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Higgs-Audio-V3 Test"
commands:
- |
timeout 20m bash -c "
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_USE_DEEP_GEMM=0
export VLLM_MOE_USE_DEEP_GEMM=0
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
"
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion · Qwen Image Edit Test"
commands:
- |
+3 -82
View File
@@ -6,13 +6,6 @@
{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
{% set default_working_dir = "/app/vllm-omni" %}
{% if env %}
env:
{% for name in env %}
{{ name }}: "{{ env[name] }}"
{% endfor %}
{% endif %}
steps:
- group: "AMD Tests"
depends_on: ~
steps:
@@ -38,15 +31,8 @@ steps:
agents:
queue: amd-cpu
{% for step_group in steps %}
{% if step_group.group %}
- group: "{{ step_group.group }}"
{% if step_group.depends_on %}
depends_on: {{ step_group.depends_on }}
{% endif %}
steps:
{% for step in step_group.steps %}
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
{% for step in steps %}
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
- label: "{{ step.agent_pool }}: {{ step.label }}"
depends_on: amd-build
agents:
@@ -60,82 +46,17 @@ steps:
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
env:
DOCKER_BUILDKIT: "1"
{% if step.env %}
{% for name in step.env %}
{{ name }}: "{{ step.env[name] }}"
{% endfor %}
{% endif %}
TEST_COMMAND: |-
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
{% if env %}
{% for name in env %}
export {{ name }}="{{ env[name] }}"
{% endfor %}
{% endif %}
{% if step.env %}
{% for name in step.env %}
export {{ name }}="{{ step.env[name] }}"
{% endfor %}
{% endif %}
{% if "mi250" in step.agent_pool %}
python3 -m pip uninstall -y amd-aiter
{% endif %}
{{ indented_cmd | safe }}
priority: 100
{% if step.timeout_in_minutes %}
timeout_in_minutes: {{ step.timeout_in_minutes }}
{% endif %}
{% if step.grade and step.grade == "Blocking" %}
soft_fail: false
{% else %}
soft_fail: true
{% endif%}
{% endif %}
{% endfor %}
{% elif step_group.mirror_hardwares and mirror_hw in step_group.mirror_hardwares %}
{% set step = step_group %}
- label: "{{ step.agent_pool }}: {{ step.label }}"
depends_on: amd-build
agents:
{% if step.agent_pool %}
queue: amd_{{ step.agent_pool }}
{% else %}
queue: amd_mi325_1
{% endif %}
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
env:
DOCKER_BUILDKIT: "1"
{% if step.env %}
{% for name in step.env %}
{{ name }}: "{{ step.env[name] }}"
{% endfor %}
{% endif %}
TEST_COMMAND: |-
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
{% if env %}
{% for name in env %}
export {{ name }}="{{ env[name] }}"
{% endfor %}
{% endif %}
{% if step.env %}
{% for name in step.env %}
export {{ name }}="{{ step.env[name] }}"
{% endfor %}
{% endif %}
{% if "mi250" in step.agent_pool %}
python3 -m pip uninstall -y amd-aiter
{% endif %}
{{ indented_cmd | safe }}
priority: 100
{% if step.timeout_in_minutes %}
timeout_in_minutes: {{ step.timeout_in_minutes }}
{% endif %}
{% if step.grade and step.grade == "Blocking" %}
soft_fail: false
{% else %}
soft_fail: true
{% endif%}
{% endif %}
{% endfor %}
{% endfor %}
-97
View File
@@ -1,97 +0,0 @@
---
name: precheck-pr
description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
---
# PR Pre-Check
Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only.
## Mode Selection
| Mode | When | Time |
|------|------|------|
| **Quick** | About to push, final sanity check | ~3 min |
| **Full** | Ready for review, want maintainer-level scan | ~10 min |
Default to quick if unsure. Run full before marking a PR "ready for review."
## Workflow
### Step 1: Detect Base Branch
```bash
BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \
|| git merge-base HEAD main 2>/dev/null \
|| echo origin/main)
echo "diffing against ${BASE_SHA}"
git diff --name-only ${BASE_SHA}...HEAD
```
### Step 2: Validate PR Title
Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes:
| Prefix | Applies to |
|--------|-----------|
| `[Bugfix]` | Bug fixes |
| `[CI/Build]` | Build or CI improvements |
| `[Doc]` | Documentation changes |
| `[Model]` | New/improved models (include model name) |
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
| `[Kernel]` | CUDA/kernel changes |
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
| `[Misc]` | Other changes (use sparingly) |
✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title.
⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add <ModelName> ...`).
### Step 3: Categorize the PR
| Diff contains | PR type |
|---------------|---------|
| New files under `vllm_omni/model_executor/models/<name>/` | **New Model** |
| Changes to `vllm_omni/diffusion/` | **Diffusion Model** |
| `[Bugfix]` prefix or single-file fix | **Bug Fix** |
| Perf/benchmark/throughput claims in commit msg or diff | **Performance** |
| Everything else | **General** |
If multiple rows apply (e.g., a diffusion model is also a new model), union the checklists.
### Step 4: Run Checklist
Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
### Step 5: Print Report
```
Pre-check report for <branch>
Mode: quick | full
Type: <new-model | diffusion-model | bug-fix | perf | general>
Dimension Result
───────────────── ──────
PR title format ✓
PR desc integrity ✓
Registry/config ✓
Dead code ⚠ 2 warnings
Accuracy ✓
Benchmark ✗ missing software versions
Verdict: 1 blocking | 2 warnings | recommend fixing ✗ before PR
```
**Severity:**
| Mark | Meaning |
|------|---------|
| ✗ | Blocking — fix before opening PR |
| ⚠ | Warning — consider fixing |
| ✓ | Pass |
| — | Skipped (not applicable) |
### Stop Here
Do not post comments, open PRs, or modify files. The report is for the contributor's terminal only.
@@ -1,135 +0,0 @@
# Pre-Check Checklists
Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing.
---
## All PRs: Title Format (check first)
PR title must follow the project convention documented in `docs/contributing/README.md`:
| Prefix | Applies to |
|--------|-----------|
| `[Bugfix]` | Bug fixes |
| `[CI/Build]` | Build or CI improvements |
| `[Doc]` | Documentation changes |
| `[Model]` | New/improved models (include model name) |
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
| `[Kernel]` | CUDA/kernel changes |
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
| `[Misc]` | Other changes (use sparingly) |
**Examples:**
```
✓ [Model] Add MiniCPM-o 4.5 support
✓ [Bugfix] Fix Ovis image text encoder dtype
✓ [Core] Qwen3-Omni performance optimization
✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend
✗ Fix Ovis image text encoder dtype ← missing prefix
✗ [bugfix] Fix Ovis image text encoder ← wrong case
✗ [WIP] Add model ← WIP in title
```
- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add <ModelName> ...`). ⚠ if missing the model name.
---
## Bug Fix PRs
### Quick
- [ ] **PR body has 4 sections:** what broke (error message), repro steps, root cause, fix description
- [ ] **Fix is minimal:** diff touches only files directly related to the bug. No unrelated refactors or style changes.
- [ ] **Repro is runnable:** the PR body includes a copy-paste command that reproduces the bug
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
### Full (adds)
- [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
- [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
- [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
- [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
- [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
---
## Performance PRs
### Quick
- [ ] **Before/after numbers in PR body:** at minimum, one metric (RTF, RPS, latency, VRAM) with main vs PR values
- [ ] **Hardware specified:** GPU model, count, VRAM. "L4" or "H20" is sufficient. "tested on GPU" is not.
- [ ] **Concurrency scaling:** if the optimization claims better throughput, numbers at 2+ concurrency levels
- [ ] **No unexplained regressions:** if any metric goes down at any concurrency level, explain why
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
### Full (adds)
- [ ] **Software versions:** torch, CUDA, vllm, vllm-omni versions in PR body
- [ ] **Warmup stated:** "1 warmup + 3 measured" or equivalent
- [ ] **Stddev or range:** not just a single best run. "average of 2 rounds" is minimum.
- [ ] **VRAM measured:** peak GPU memory before and after
- [ ] **Benchmark script checked in:** under `tests/` or `examples/`, runnable by reviewer
- [ ] **Exact command line in PR:** not pseudocode, not "run the benchmark"
- [ ] **Pytest summary included:** `36 passed, 3 skipped, 17 warnings in 2897.66s` — proves the whole suite passes
- [ ] **Config changes documented:** if `max_num_seqs`, `gpu_memory_utilization`, or deploy YAML changed, old and new values stated with reasoning
---
## New Model PRs
### Quick
- [ ] **PR body file list matches `git diff --name-only`:** every file path in the description exists in the diff
- [ ] **Registry entries match claimed count:** count `_OMNI_MODELS` or `_DIFFUSION_MODELS` entries added vs claimed in body
- [ ] **Pipeline config `custom_process_input_func` paths resolve:** each referenced module.function exists
- [ ] **`__init__.py` exports match registry:** each registered class appears in `__init__.py` and `__all__`
- [ ] **Shell script modes match README:** `run_curl.sh` supports every mode the README documents
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
### Full (adds)
- [ ] **Dead code scan (5 patterns):**
1. Dead `forward()` — defined but never called through the pipeline path
2. Dead factory/builder functions — instantiated nowhere
3. Dead wrappers — thin wrappers with no callers
4. Dead branch guards — `if version > X` where X is always satisfied
5. Unused parameters — `__init__` args never accessed
- [ ] **Copy-paste detection:** no string constant defined in 3+ files, no cross-module validation duplication, no near-identical shape coercion fns
- [ ] **Import hygiene:** no re-export-only imports, no module-level side effects without docs, no `import os` in function body when already at top. Function-body imports are acceptable when the import is only reachable from a closure or factory return value (i.e., `import` would otherwise execute at module load even though the user may never trigger that path). Either lift to top-of-file or annotate with `# noqa: PLC0415`.
- [ ] **Accuracy:** at least one test compares output against reference; audio = valid WAV at correct sample rate; determinism verified if seed claimed. ⚠ if missing and not explicitly deferred to a follow-up PR or RFC section; — if explicitly deferred (the deferral target SHOULD be linked in the PR body).
- [ ] **Performance:** hardware, software versions, warmup, and metrics (RTF, VRAM, RPS, latency) stated. ⚠ if missing and not explicitly deferred; — if explicitly deferred.
- [ ] **Benchmark settings:** model config (TP, PP, max_model_len, enforce_eager, quant), runtime config (batch, input/output spec), environment versions
- [ ] **Benchmark script checked in + exact command line + pytest summary**
---
## Diffusion Model PRs
Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model checklist above, verify:
- [ ] **Transformer adapter follows contract:** inherits from `nn.Module`, implements `load_weights()`, uses `vllm_omni.diffusion.attention.layer.Attention`
- [ ] **Pipeline has both offline and online paths:** `forward()` for batch generation, streaming path for online serving
- [ ] **Cache-DiT integration:** at least one acceleration (TeaCache, FBCache, etc.) is wired up. ✗ for new **top-level** pipelines. For **subclass** pipelines that inherit `__call__` from a parent where Cache-DiT is already absent, downgrade to ⚠ and recommend filing a follow-up issue against the parent pipeline family.
- [ ] **Parallelism config:** TP/SP/USP/CFG-Parallel options exposed in pipeline config
- [ ] **Docs table updated:** model added to `docs/models/supported_models.md` with correct metadata
- [ ] **E2E test exists:** at least one test exercises the full generate path
---
## General PRs
### Quick
- [ ] **PR body matches diff:** description doesn't claim changes not in the diff
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
- [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists
### Full (adds)
- [ ] **No new dead code in changed files:** scan for unreachable paths, unused imports
- [ ] **Test coverage:** if core code changed (`engine/`, `stages/`, `connectors/`), tests are added or existing tests cover the path
- [ ] **Import hygiene:** no redundant imports added
- [ ] **No unrelated changes:** diff scope matches PR description
-5
View File
@@ -25,12 +25,7 @@ include:
- `add-tts-model`: covers integration of new TTS models and related serving
workflows
- `generate-release-note`: helps prepare release notes for repository changes
- `precheck-pr`: self-check a branch before creating a PR — validates PR title
format, catches dead code, verifies accuracy/perf claims, and confirms merge
readiness
- `review-pr`: provides a structured workflow for reviewing pull requests
- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
latest vllm-ascend NPUModelRunner
## Maintenance Guidelines
+13 -7
View File
@@ -1,15 +1,21 @@
<!-- markdownlint-disable -->
PLEASE FILL IN THE PR DESCRIPTION HERE.
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
## Purpose
## Test Plan
**vLLM Version:**
**vLLM-Omni Commit:**
## Test Result
**BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)
---
<details>
<summary> Essential Elements of an Effective PR Description Checklist </summary>
- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
</details>
**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
-2
View File
@@ -1,5 +1,3 @@
# Contributing to vLLM-Omni
You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
+2 -4
View File
@@ -16,7 +16,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
---
*Latest News* 🔥
- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
@@ -57,9 +56,8 @@ vLLM-Omni is flexible and easy to use with:
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
- Omni-modality models (e.g. Qwen-Omni)
- Multi-modality generation models (e.g. Qwen-Image)
## Getting Started
-9
View File
@@ -151,9 +151,6 @@ class VllmOmniImageClient:
seed: int | None = None,
negative_prompt: str | None = None,
output_compression: int | None = None,
bot_task: str | None = None,
sys_type: str | None = None,
system_prompt: str | None = None,
) -> Image.Image:
if not isinstance(images, list):
images = [images]
@@ -173,12 +170,6 @@ class VllmOmniImageClient:
data["negative_prompt"] = negative_prompt
if output_compression is not None:
data["output_compression"] = str(output_compression)
if bot_task is not None:
data["bot_task"] = bot_task
if sys_type is not None:
data["sys_type"] = sys_type
if system_prompt is not None:
data["system_prompt"] = system_prompt
files = [
(
@@ -64,47 +64,6 @@ def resolve_model_name(*, model_name: str | None, model: str | None = None, outp
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
def _coerce_score_int(value: Any) -> int:
if isinstance(value, bool):
raise TypeError(f"Expected numeric score, got bool: {value!r}")
if isinstance(value, (int, float)):
return int(value)
if isinstance(value, str):
stripped = value.strip()
if stripped.lstrip("-").isdigit():
return int(stripped)
raise TypeError(f"Expected numeric score, got {type(value).__name__}: {value!r}")
def _expand_score_entry(entry: Any) -> list[int]:
if isinstance(entry, dict):
if "score" in entry:
return [_coerce_score_int(entry["score"])]
if "value" in entry:
return [_coerce_score_int(entry["value"])]
numeric_values = [_coerce_score_int(value) for value in entry.values()]
if numeric_values:
return numeric_values
raise TypeError(f"Score dict has no numeric values: {entry!r}")
return [_coerce_score_int(entry)]
def _normalize_score_list(score: Any) -> list[int]:
if isinstance(score, dict):
entries: list[Any] = [score]
elif isinstance(score, list):
entries = score
else:
entries = [score]
normalized: list[int] = []
for entry in entries:
normalized.extend(_expand_score_entry(entry))
if not normalized:
raise ValueError(f"Empty score list after normalization: {score!r}")
return normalized
def parse_score_payload(raw_text: str) -> dict[str, Any]:
try:
parsed = extract_json_object(raw_text)
@@ -114,7 +73,10 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]:
parsed = {"score": json.loads(stripped), "reasoning": ""}
else:
raise
parsed["score"] = _normalize_score_list(parsed.get("score", []))
score = parsed.get("score", [])
if not isinstance(score, list):
score = [score]
parsed["score"] = [int(value) for value in score]
return parsed
@@ -429,9 +391,6 @@ class GEditBenchRunner:
num_inference_steps: int = 20,
guidance_scale: float | None = None,
seed: int | None = 42,
bot_task: str | None = None,
sys_type: str | None = None,
system_prompt: str | None = None,
):
self.dataset_ref = dataset_ref
self.output_root = output_root
@@ -441,9 +400,6 @@ class GEditBenchRunner:
self.num_inference_steps = num_inference_steps
self.guidance_scale = guidance_scale
self.seed = seed
self.bot_task = bot_task
self.sys_type = sys_type
self.system_prompt = system_prompt
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
def generate(
@@ -522,9 +478,6 @@ class GEditBenchRunner:
num_inference_steps=self.num_inference_steps,
guidance_scale=self.guidance_scale,
seed=self.seed,
bot_task=self.bot_task,
sys_type=self.sys_type,
system_prompt=self.system_prompt,
)
edited_image.save(output_path)
return {
@@ -680,25 +633,6 @@ def build_parser() -> argparse.ArgumentParser:
generate.add_argument("--num-inference-steps", type=int, default=20)
generate.add_argument("--guidance-scale", type=float, default=None)
generate.add_argument("--seed", type=int, default=42)
generate.add_argument(
"--bot-task",
type=str,
default=None,
choices=["think", "recaption", "think_recaption", "vanilla"],
help="HunyuanImage-3 prompting mode for /v1/images/edits (optional).",
)
generate.add_argument(
"--sys-type",
type=str,
default=None,
help="Hunyuan system-prompt type (e.g. en_unified). Defaults from --bot-task when omitted.",
)
generate.add_argument(
"--system-prompt",
type=str,
default=None,
help="Custom Hunyuan system prompt body (sys_type=custom).",
)
generate.add_argument("--workers", type=int, default=1)
generate.add_argument("--max-samples", type=int, default=None)
generate.add_argument("--samples-per-group", type=int, default=None)
@@ -741,9 +675,6 @@ def main(argv: list[str] | None = None) -> int:
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
bot_task=args.bot_task,
sys_type=args.sys_type,
system_prompt=args.system_prompt,
)
records = runner.generate(
model_name=model_name,
@@ -1,10 +0,0 @@
se001|ocean waves crashing on a rocky beach|10.0
se002|heavy rain on a metal roof at night|8.0
se003|a dog barking in the distance|6.0
se004|a crackling campfire in a quiet forest|10.0
se005|footsteps on dry leaves|6.0
se006|thunder rolling across a wide valley|8.0
se007|a busy coffee shop with chatter and espresso machines|10.0
se008|wind howling through tall pine trees|8.0
se009|a small stream flowing over pebbles|10.0
se010|a creaking wooden door slowly opening|6.0
@@ -1,10 +0,0 @@
ttsd001|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Hello and welcome to the show today. [S2] Thanks for having me, it's a pleasure to be here. [S1] Let's dive right in.
ttsd002|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you finish the report? [S2] Almost — I'm just polishing the conclusion. [S1] Great, send it over when you're done.
ttsd003|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have you tried that new restaurant downtown? [S2] Yes, the pasta is amazing. [S1] I'll book a table for Friday.
ttsd004|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] What did you think of yesterday's meeting? [S2] Productive, but a bit long. [S1] Agreed, we should set an hour cap next time.
ttsd005|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] The weather is beautiful today. [S2] Perfect for a hike. [S1] Want to drive out to the trails this afternoon?
ttsd006|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] How did the demo go? [S2] Better than expected — the client signed off on the prototype. [S1] Fantastic news!
ttsd007|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Coffee or tea this morning? [S2] Coffee please, with a splash of milk. [S1] Coming right up.
ttsd008|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you read the email about the deadline change? [S2] I did, we now have until next Friday. [S1] That should give us enough buffer.
ttsd009|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Are you joining the team dinner tonight? [S2] I'll be there around seven. [S1] See you then.
ttsd010|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have we backed up the database? [S2] Last backup ran two hours ago. [S1] Good, I'll trigger a fresh snapshot before the migration.
View File
Whitespace-only changes.
-16
View File
@@ -53,8 +53,6 @@ _TASK_TO_DATASET: dict[str, str] = {
"voice_clone": "seed-tts",
"default_voice": "seed-tts-text",
"voice_design": "seed-tts-design",
"dialogue": "ttsd",
"sound_effect": "sound-effect",
}
# Default design dataset path (bundled with the repo)
@@ -83,7 +81,6 @@ def build_bench_args(
output_dir: str | None,
result_filename: str | None,
extra_cli_args: list[str],
output_len: int | None = None,
) -> list[str]:
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
dataset_name = _TASK_TO_DATASET[task]
@@ -142,9 +139,6 @@ def build_bench_args(
if wer_eval:
cmd.append("--seed-tts-wer-eval")
if output_len is not None:
cmd += ["--hf-output-len", str(output_len)]
if output_dir or result_filename:
out_dir = output_dir or "."
os.makedirs(out_dir, exist_ok=True)
@@ -244,15 +238,6 @@ def main() -> None:
parser.add_argument("--host", default="localhost")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
parser.add_argument(
"--output-len",
type=int,
default=None,
help="Override per-request max_new_tokens forwarded to the server "
"(seed-tts datasets default to 2048; the MOSS-TTS full talker takes "
"this as 'max_new_frames' so smaller values cap audio length per "
"request and keep bench runtime tractable).",
)
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
args = parser.parse_args()
@@ -307,7 +292,6 @@ def main() -> None:
output_dir=args.output_dir,
result_filename=result_filename,
extra_cli_args=args.extra or [],
output_len=args.output_len,
)
result = run_one_benchmark(cmd)
if result is not None:
-61
View File
@@ -14,7 +14,6 @@ models:
supported_tasks: [default_voice, voice_design]
backend: openai-audio-speech
endpoint: /v1/audio/speech
trust_remote_code: true
task_extra_body:
default_voice:
voice: Vivian
@@ -28,7 +27,6 @@ models:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
trust_remote_code: true
task_extra_body:
voice_clone:
task_type: Base
@@ -38,66 +36,7 @@ models:
trust_remote_code: true
backend: openai-audio-speech
endpoint: /v1/audio/speech
trust_remote_code: true
task_extra_body:
voice_clone: {}
default_voice:
voice: default
# MOSS-TTS family. After the serving_speech.py variant-detection refactor
# all five variants now go through OpenAI /v1/audio/speech with the
# variant-specific request fields (ambient_sound, ref_audio_2, instructions).
OpenMOSS-Team/MOSS-TTS:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone: {}
# MOSS-TTS-v1.5: continued-training upgrade of MOSS-TTS 1.0 — same
# MossTTSDelay-8B architecture and API, so identical wiring. Adds 31-language
# synthesis via the `language` field and inline `[pause Xs]` markers.
OpenMOSS-Team/MOSS-TTS-v1.5:
supported_tasks: [voice_clone]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone: {}
OpenMOSS-Team/MOSS-TTS-Realtime:
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.
OpenMOSS-Team/MOSS-TTSD-v1.0:
supported_tasks: [voice_clone, dialogue]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_clone: {}
dialogue: {}
# MOSS-SoundEffect ignores text/ref_audio; expects ambient_sound +
# duration_seconds. The sound_effect dataset module emits those.
OpenMOSS-Team/MOSS-SoundEffect:
supported_tasks: [sound_effect]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
sound_effect:
duration_seconds: 10
# MOSS-VoiceGenerator: instructions-driven (no ref_audio). seed-tts-design
# dataset already emits {instructions, target_text}, so voice_design works
# out of the box once the serving validator accepts instructions.
OpenMOSS-Team/MOSS-VoiceGenerator:
supported_tasks: [voice_design]
backend: openai-audio-speech
endpoint: /v1/audio/speech
task_extra_body:
voice_design: {}
+1 -4
View File
@@ -1,8 +1,5 @@
# Target vllm version: v0.21.0 (releases/v0.21.0 branch).
# 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.21.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
ARG BASE_IMAGE=vllm/vllm-openai:v0.21.0
FROM ${BASE_IMAGE}
ARG COMMON_WORKDIR=/app
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
+1 -1
View File
@@ -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.21.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
+1
View File
@@ -92,6 +92,7 @@ Core scheduling and caching components.
Configuration classes.
- [vllm_omni.config.model.OmniModelConfig][]
- [vllm_omni.config.stage_config.ModelPipeline][]
- [vllm_omni.config.stage_config.StageConfig][]
- [vllm_omni.config.stage_config.StageConfigFactory][]
- [vllm_omni.config.stage_config.StageType][]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 403 KiB

After

Width:  |  Height:  |  Size: 408 KiB

+26 -134
View File
@@ -2,29 +2,29 @@
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
## Notes
We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend.
We currently only support using the "openai-chat-omni" backend.
## Basic Parameter Description
You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
- `--omni`
- `--omni`
Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
- `--backend`
- `--backend`
Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
- `--model`
- `--model`
The model identifier to load, filled according to the models supported by vLLM-Omni.
- `--endpoint`
- `--endpoint`
The API endpoint exposed externally, to which clients send their requests.
- `--dataset-name`
- `--dataset-name`
The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
- `--num-prompts`
- `--num-prompts`
The total number of requests to send, an integer.
- `--max-concurrency`
- `--max-concurrency`
"Maximum number of concurrent requests. This can be used "
"to help simulate an environment where a higher level component "
"is enforcing a maximum number of concurrent requests. While the "
@@ -34,73 +34,69 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
"actual request rate may be lower than specified with --request-rate, "
"if the server is not processing requests fast enough to keep up."
- `--request-rate`
- `--request-rate`
"Number of requests per second. If this is inf, "
"then all the requests are sent at time 0. "
"Otherwise, we use Poisson process or gamma distribution "
"to synthesize the request arrival times."
- `--ignore-eos`
- `--ignore-eos`
"Set ignore_eos flag when sending the benchmark request."
- `--metric-percentiles`
- `--metric-percentiles`
Comma-separated list of percentiles for selected metrics. "
"To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
"Default value is \"99\"."
"Use \"--percentile-metrics\" to select metrics.
- `--percentile-metrics`
- `--percentile-metrics`
"Comma-separated list of selected metrics to report percentiles."
"This argument specifies the metrics to report percentiles."
'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", '
'"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
- `-print-stage`
Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
- `--save-result`
- `--save-result`
Specify to save benchmark results to a json file
- `--save-detailed`
- `--save-detailed`
"When saving the results, whether to include per request "
"information such as response, error, ttfs, tpots, etc."
- `--result-dir`
- `--result-dir`
"Specify directory to save benchmark json results."
"If not specified, results are saved in the current directory."
- `--result-filename`
- `--result-filename`
"Specify the filename to save benchmark json results."
"If not specified, results will be saved in "
"{label}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
- `--random-prefix-len`
- `--random-prefix-len`
Number of fixed prefix tokens before the random context in a request.
The total input length is the sum of random-prefix-len and a random
context length sampled from [input_len * (1 - range_ratio),
input_len * (1 + range_ratio)].Only the random and random-mm modes
support this parameter.
- `--random-input-len`
- `--random-input-len`
Number of input tokens per request.Only the random and random-mm modes support this parameter.
- `--random-output-len`
- `--random-output-len`
Number of output tokens per request.Only the random and random-mm modes support this parameter.
- `--random-range-ratio`
- `--random-range-ratio`
Range ratio for sampling input/output length,
used only for random sampling. Must be in the range [0, 1) to define
a symmetric sampling range
[length * (1 - range_ratio), length * (1 + range_ratio)].
Only the random and random-mm modes support this parameter.
- `--random-mm-base-items-per-request`
- `--random-mm-base-items-per-request`
Base number of multimodal items per request for random-mm.
Actual per-request count is sampled around this base using
--random-mm-num-mm-items-range-ratio.
Only the random-mm mode supports this parameter.
- `--random-mm-limit-mm-per-prompt`
- `--random-mm-limit-mm-per-prompt`
Per-modality hard caps for items attached per request, e.g.
'{"image": 3, "video": 1, "audio": 1}'. The sampled per-request item
count is clamped to the sum of these limits. When a modality
@@ -108,7 +104,7 @@ Specify to save benchmark results to a json file
renormalized.
Only the random-mm mode supports this parameter.
- `--random-mm-num-mm-items-range-ratio`
- `--random-mm-num-mm-items-range-ratio`
Range ratio r in [0, 1] for sampling items per request.
We sample uniformly from the closed integer range
[floor(n*(1-r)), ceil(n*(1+r))]
@@ -119,7 +115,7 @@ Specify to save benchmark results to a json file
An error is raised if the computed min exceeds the max.
Only the random-mm mode supports this parameter.
- `--random-mm-bucket-config`
- `--random-mm-bucket-config`
The bucket config is a dictionary mapping a multimodal item
sampling configuration to a probability.
Currently allows for 3 modalities: audio, images and videos.
@@ -280,9 +276,7 @@ Generate synthetic image、video、audio inputs alongside random text prompts to
Notes:
- Works only with online benchmark via:
- the OpenAI chat backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
- the OpenAI edit image backend (`--backend openai-image-edits-omni`) and endpoint `/v1/images/edits`.
- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
Start the server (example):
@@ -363,105 +357,3 @@ How sampling works:
This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
</details>
### Multi-Stage Benchmark
<details class="admonition abstract" markdown="1">
<summary>Show more</summary>
Start the server (change the port and model):
```bash
vllm serve --omni --port 29999 --model ~/Qwen3-Omni-30B-A3B-Instruct
```
Then run the benchmarking script (remember to add the flag `--print-stage`):
```bash
vllm bench serve --omni \
--dataset-name random \
--port 29999 \
--max-concurrency 1 \
--num-warmups 2 \
--model ~/Qwen3-Omni-30B-A3B-Instruct \
--endpoint /v1/chat/completions \
--backend openai-chat-omni \
--num-prompts 3 \
--random-input-len 2500 \
--ignore-eos \
--percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \
--random-output-len 900 \
--extra_body '{"modalities": ["text","audio"]}' \
--print-stage
```
Besides the "Serving Benchmark Result", there will be "Stage Benchmark Result" below:
```text
============= Stage Benchmark Result =============
=============== Stage 0 (thinker) ================
-------------------Stage Timing-------------------
Mean stage_gen_time (ms): 18243.16
Median stage_gen_time (ms): 18220.83
P99 stage_gen_time (ms): 18296.75
================== Text Result ===================
Stage generated tokens: 2700
-----------Serving Time to First Token------------
Mean Serving TTFT (ms): 816.31
Median Serving TTFT (ms): 818.94
P99 Serving TTFT (ms): 829.76
-----Time per Output Token (excl. 1st token)------
Mean TPOP (ms): 19.41
Median TPOP (ms): 19.40
P99 TPOP (ms): 19.47
---------------Inter-token Latency----------------
Mean ITL (ms): 19.41
Median ITL (ms): 19.12
P99 ITL (ms): 23.58
================ Stage 1 (talker) ================
-------------------Stage Timing-------------------
Mean stage_gen_time (ms): 18849.64
Median stage_gen_time (ms): 13641.15
P99 stage_gen_time (ms): 33655.58
============= Internal Stream Result =============
-----------Serving Time to First Chunk------------
Mean Serving TTFC (ms): 1029.14
Median Serving TTFC (ms): 1027.98
P99 Serving TTFC (ms): 1032.59
-----Time per Output Chunk (excl. 1st chunk)------
Mean TPOP (ms): 17.53
Median TPOP (ms): 18.70
P99 TPOP (ms): 18.88
---------------Inter-chunk Latency----------------
Mean ICL (ms): 16.27
Median ICL (ms): 18.19
P99 ICL (ms): 25.86
=============== Stage 2 (code2wav) ===============
-------------------Stage Timing-------------------
Mean stage_gen_time (ms): 18860.74
Median stage_gen_time (ms): 13655.84
P99 stage_gen_time (ms): 33659.79
================== Audio Result ==================
Stage audio duration generated(s): 260.11
Stage audio frames generated: 6242730
-----------Serving Time to First Packet-----------
Mean Serving AUDIO_TTFP (ms): 1366.16
Median Serving AUDIO_TTFP (ms): 1367.20
P99 Serving AUDIO_TTFP (ms): 1368.26
-----------------Real Time Factor-----------------
Mean AUDIO_RTF: 0.24
Median AUDIO_RTF: 0.26
P99 AUDIO_RTF: 0.27
```
Explanation:
- stage_gen_time: Time from submitting a request to a specific stage to that stage finishing generation.
- Serving TTFC (Time to First Chunk): Time from the HTTP request being accepted by the serving frontend to
the stage producing its first non-empty output chunk. Keep the name TTFT in text stage, TTFP in audio stage
for easier mapping.
- TPOP (Time per Output Chunk): Average time from the first output chunk to stage completion, divided by the
number of remaining output chunks. The TPOP abbreviation follows Qwen3.5-Omni Technical Report.
- ICL (Inter-Chunk Latency): Time between two consecutive output chunks produced by the same stage. Keep the
name ITL in text stage for easier mapping.
</details>
-10
View File
@@ -116,16 +116,6 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
!!! note
If the PR spans more than one category, please include all relevant prefixes.
### Pre-Check Before Submitting
Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-review against project conventions:
The skill offers two modes:
- **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
- **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene
The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) for the full checklist.
### Local Test
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
@@ -8,7 +8,6 @@ Backend-specific setup lives in separate docs:
- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
- [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md)
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
@@ -24,7 +23,6 @@ Current connectors operate in D2H2D (device to host to device) mode.
| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
| Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. |
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
@@ -1,76 +0,0 @@
# MoriTransferEngineConnector
## When to Use
Currently supports intra-node deployment with Mori.
As noted in #1742, inter-node support will be added back in a future
refactor.
## Mechanism
Uses Mori's `IOEngine` / `MemoryDesc` API for zero-copy RDMA transfers.
- Data Plane: RDMA (InfiniBand/RoCE) with managed memory pool.
- Control Plane: ZMQ for pull-request handshake and async completion.
## Installation
See the [Mori repository](https://github.com/ROCm/mori) for installation instructions.
## Configuration
Mori is configured through the new deploy-config schema (see
[`docs/configuration/stage_configs.md`](../../../configuration/stage_configs.md)).
Define the connector at the top level of the deploy YAML and reference it
by name from each stage's `input_connectors` / `output_connectors`:
```yaml
connectors:
mori_connector:
name: MoriTransferEngineConnector
extra:
host: "auto"
zmq_port: 50051
device_name: ""
memory_pool_size: 536870912
memory_pool_device: "cuda"
stages:
- stage_id: 0
output_connectors:
to_stage_1: mori_connector
- stage_id: 1
input_connectors:
from_stage_0: mori_connector
```
A ready-to-run intra-node example for Qwen3-Omni-MoE on AMD MI300X lives
at `vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml` and can be loaded
with:
```bash
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \
--deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml
```
The yaml wires `MoriTransferEngineConnector` (with `backend_type: xgmi`) to
the chunk_transfer_adapter path (`async_chunk: true`) so stage-to-stage
hidden-state and codec-frame streams ship GPU-to-GPU over AMD Infinity
Fabric instead of SHM. Qwen2.5-Omni + Mori is not yet functional on
the chunk path: it needs `thinker2talker_async_chunk` /
`talker2code2wav_async_chunk` input processors that do not exist yet
(the orchestrator-level path the upstream PR originally targeted was
lost during the entrypoints → engine refactor and removed in #1742).
Parameters:
- host: local RDMA IP (`"auto"` for auto-detect).
- zmq_port: ZMQ base port for control-plane communication.
- device_name: RDMA device (e.g., `"mlx5_0"`), empty for auto-detect.
- memory_pool_size: RDMA memory pool size in bytes.
- memory_pool_device: `"cpu"` (pinned) or `"cuda"` (GPUDirect / XGMI RDMA).
For more details, refer to the
[Mori repository](https://github.com/ROCm/mori).
-4
View File
@@ -13,10 +13,6 @@ This section contains design documents and architecture specifications for vLLM-
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
## Infrastructure Design Documents
- [Prometheus Metrics](metrics.md)
## Module Design Documents
- [AR Module](module/ar_module.md)
-208
View File
@@ -1,208 +0,0 @@
# Prometheus Metrics Design
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
## Objectives
- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
## Background
### Upstream vLLM Metrics
Upstream vLLM defines ~37 Prometheus metric families under the `vllm:` prefix. These are registered by `PrometheusStatLogger` and cover engine-level state: KV cache usage, running/waiting request counts, token throughput, TTFT, inter-token latency, e2e latency, and so on. They are served via the `/metrics` HTTP endpoint provided by `prometheus_fastapi_instrumentator` and the default `prometheus_client` WSGI handler.
vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` collector whose `_name` attribute contains the substring `"vllm"`. This runs during engine initialization to clean up stale collectors from prior instantiations within the same process.
### The Problem
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
## Architecture
### Component Overview
```text
+------------------------+
| API Server (FastAPI) |
| GET /metrics |
+-----------+------------+
|
prometheus_client default registry
|
+--------+--------+--------+--------+--------+
| |
vllm:omni_* vllm:*
collectors collectors
| |
+----+--------+ +-----------+ +----------+ +-----------+
| OmniPromet- | | OmniMod- | | OmniTra- | | OmniProm- |
| heusMetrics | | alityMet- | | nsferMe- | | etheusSt- |
| | | rics | | trics | | atLogger |
+----+--------+ +-----+-----+ +----+-----+ +----+------+
| | | |
OmniBase OmniBase Orchestrator Orchestrator
(request life- (finalize + (record_trans- (per-(stage,
cycle, success/ streaming fer_tx/rx replica)
fail counter) hooks via hooks via scheduler/
observe_* emit hook in iteration
APIs) OrchestratorAg- stats)
gregator)
```
### Data Flow
There are four independent paths for metric collection.
**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
- `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path).
- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transfer_rx`. After the existing `TransferEdgeStats` accumulation, the aggregator calls `_emit_transfer_tx` / `_emit_transfer_rx`. Those:
1. Resolve `from_replica` / `to_replica` via a `replica_resolver` callback supplied by `async_omni.py`. The resolver delegates to `stage_pool.get_bound_replica_id(request_id)` — i.e. the orchestrator's existing sticky-routing binding is the source of truth.
2. Convert the underlying `_ms` accumulators to seconds and call the `_s`-suffixed observe methods on `OmniTransferMetrics`.
Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
### Shared State Between Threads
The Orchestrator runs in a background thread. The API server (OmniBase) runs in the asyncio event loop thread. `OmniRequestCounter` bridges them — a plain Python object with an `int` field. The Orchestrator increments/decrements it; the entrypoint reads it for gauge updates. No lock is needed because the counter is advisory (a stale read by one Prometheus scrape interval is acceptable). It is created by `AsyncOmniEngine.__init__()` and passed to the Orchestrator at construction time.
### Metric Registration and Lifecycle
All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
## OmniPrometheusStatLogger Wrap
Upstream `PrometheusStatLogger.__init__` hard-codes `labelnames = ["model_name", "engine"]` as a local variable, references it across ~37 metric-family construction sites, and uses the `engine` label value in five different `.labels()` call shapes (kwarg with int engine, kwarg with str engine, positional with str engine in the middle, plus a `metrics_info["engine"] = str(...)` dict pattern). To reshape `engine` into `stage` + `replica` without forking the entire upstream `__init__`, the wrap uses three coordinated mechanisms:
1. **Class-level metric class slot overrides.** `OmniPrometheusStatLogger` overrides `_gauge_cls`, `_counter_cls`, `_histogram_cls` (which upstream calls via `self._gauge_cls(...)` etc.) with `_RelabelGauge` / `_RelabelCounter` / `_RelabelHistogram` wrapper classes. These intercept the `labelnames` kwarg at metric family creation time and replace `engine` with `("stage", "replica")`.
2. **Property descriptor for `per_engine_labelvalues`.** Upstream builds `self.per_engine_labelvalues = {idx: [model_name, str(idx)]}` inside `__init__` and then captures it into a local variable for `create_metric_per_engine` calls. By making `per_engine_labelvalues` a Python property on the subclass, the setter intercepts upstream's assignment and rewrites each 2-tuple into a 3-tuple `[model_name, stage, replica]` using the `stage_replica_map` supplied at construction time. The captured local then sees the rewritten dict.
3. **Override of `.labels()` on the wrapper classes.** For the five call sites that pass `engine` directly (kwarg or positional, int or str), `_RelabelMixin.labels()` translates the engine value back to `(stage, replica)` via a process-level `_ENGINE_INDEX_MAP` populated by `OmniPrometheusStatLogger.__init__`. This handles `gauge_engine_sleep_state.labels(engine=idx, ...)`, `counter_request_success_base.labels(model_name, str(idx), str(reason))`, `info_gauge.labels(**metrics_info)`, etc.
The three sub-helpers that upstream `PrometheusStatLogger.__init__` constructs (`spec_decoding_prom` / `kv_connector_prom` / `perf_metrics_prom`) use their own `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots and would otherwise build families with the raw 2-element labelnames. `_OmniPerfMetricsProm` / `_OmniSpecDecodingProm` / `_OmniKVConnectorProm` subclass each helper to route the same relabel mixin through their internal family construction.
The `Orchestrator` constructs `stage_replica_map` from the static `stage_pools` configuration at startup:
```python
stage_replica_map = {
flat_idx: (str(stage_id), str(replica_id))
for flat_idx, (stage_id, replica_id) in enumerate(
(s, r)
for s, pool in enumerate(stage_pools)
for r in range(pool.num_replicas)
)
}
```
A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestrator so the per-replica `record(engine_idx=...)` call site can look up the right flat index.
> Dynamic add/remove of replicas at runtime is intentionally out of scope — the upstream `PrometheusStatLogger` materializes per-engine_idx child metrics at init time, and supporting hot-add would require non-trivial intervention into upstream's per-family child dictionaries.
## Gating: `--log-stats`
All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
Behavior with `--log-stats=off` (default):
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
Behavior with `--log-stats=on`: all metric paths fire normally; the orchestrator's per-replica recording is bounded only by `OmniSchedulerMixin.make_stats()`'s per-scheduler 1 Hz throttle (see next section).
The overhead with the flag on is small enough that an A/B benchmark on Qwen3-Omni-30B single replica (30 sequential audio requests) showed a mean latency delta of +0.6% (Welch's t = 0.318, n=30, not statistically significant at α=0.05); the `/metrics` line count drops from 1358 to 124 lines when the flag is off.
## Throttling: `make_stats()` Override
Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, returning a SchedulerStats object for the orchestrator. Under vLLM's architecture, this is fine. But since vLLM-Omni requires that the object be serialized and transferred over ZMQ, receiving a SchedulerStats object on every step can introduce unacceptable overhead to the system.
`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
## Metric Definitions
### Pipeline (4)
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
### Audio (7)
Labels: `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description |
|--------|------|-------------|-------------|
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
### Cross-stage transfer (4)
Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
### LLM stage-level (wrapped `vllm:*`)
After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — carries `{model_name, stage, replica}` labels. For the full upstream catalog see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md); note that metrics depending on features unsupported in vLLM-Omni (e.g. speculative decoding, LoRA) will not be available.
## Naming Convention
- All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used:
- `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values.
- `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix.
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
## Logging vs. Prometheus
`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the `INFO` log. This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing.
`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics` form the Prometheus-oriented path. They record aggregate counters, gauges, and histograms suitable for time-series monitoring and alerting. Both paths share the same source data (`StageRequestStats`, `TransferEdgeStats`) — `OrchestratorAggregator.record_transfer_tx/rx` in particular calls both the existing accumulator code and the Prometheus emit hook in the same method body. The two consumption models can run simultaneously without coupling.
The separation follows upstream vLLM's pattern of `LoggingStatLogger` vs. `PrometheusStatLogger` — same underlying data, different consumption models.
+2 -2
View File
@@ -305,8 +305,8 @@ class GPUWorker:
)
# Load model
model_loader = DiffusersPipelineLoader(load_config, self.od_config)
self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}")
model_loader = DiffusersPipelineLoader(load_config)
self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}")
# Setup cache backend
from vllm_omni.diffusion.cache.selector import get_cache_backend
@@ -411,8 +411,6 @@ Notes:
- `runtime.max_batch_size` controls stage-level batching.
- Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
- Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
#### 2) Enable async chunk
+1 -1
View File
@@ -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, and 0.21.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"
@@ -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.21.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.21.0 release version of the library
```bash
uv pip install vllm==0.22.0 --torch-backend=auto
uv pip install vllm==0.21.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.21.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.21.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.21.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_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.21.0 \
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -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.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.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.21.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.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.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.21.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.21.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.21.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.21.0
```
# --8<-- [end:pre-built-images]
+3 -3
View File
@@ -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.21.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.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
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.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
## Offline Inference
+1 -2
View File
@@ -119,10 +119,9 @@ def extract_omni_serve_subparser_init():
exec_globals = {
"_FlexibleArgumentParser": _FlexibleArgumentParser,
"FlexibleArgumentParser": _FlexibleArgumentParser,
"TrackingArgumentParser": _FlexibleArgumentParser, # only needs add_argument, not tracking
"TrackingNamespace": None, # type annotation only
"make_arg_parser": lambda parser: parser, # no-op for doc
"_ensure_vllm_platform": lambda: None, # no-op for doc
"nullify_stage_engine_defaults": lambda parser: None, # no-op for doc
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
-11
View File
@@ -33,7 +33,6 @@ 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
@@ -42,7 +41,6 @@ th {
| `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` | ✅︎ | | | |
| `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` | ✅︎ | ✅︎ | | ✅︎ |
@@ -50,8 +48,6 @@ th {
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
@@ -63,12 +59,7 @@ th {
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
@@ -76,13 +67,11 @@ th {
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
+1 -39
View File
@@ -116,7 +116,7 @@ Content-Type: application/json
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
**Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
#### Voice Clone Parameters (Base task)
@@ -381,44 +381,6 @@ are cached in-process with a shared LRU so repeated requests with the same
all TTS model types; deleting a voice invalidates every model-type slot at
once.
### Precomputed Custom Voices
Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
Generate a directory containing `custom_voice_manifest.json` plus one
`.safetensors` file per voice, then set the pipeline-wide deploy config field:
```yaml
custom_voice_dir: /path/to/custom_voices
```
Qwen3-TTS profiles are created with:
```bash
python examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio" \
--mode icl \
--output-dir /path/to/custom_voices
```
VoxCPM2 profiles are created with:
```bash
python examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--mode ref_continuation \
--prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices
```
Only profiles whose safetensors payload can be loaded and validated are exposed
by `GET /v1/audio/voices`. Valid precomputed voices can be used in
`POST /v1/audio/speech` by passing `voice="alice"` without `ref_audio`.
**Configuration (environment variables):**
| Variable | Default | Description |
+3 -44
View File
@@ -67,9 +67,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
| `width` | integer | model default | Output video width |
| `height` | integer | model default | Output video height |
| `num_frames` | integer | model default | Number of generated frames |
@@ -81,8 +80,6 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
| `flow_shift` | number | null | Scheduler flow-shift value |
| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
| `seed` | integer | null | Random seed for reproducibility |
| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation |
| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration |
| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
@@ -126,41 +123,6 @@ curl -s http://localhost:8091/v1/videos \
-F "fps=16"
```
### Video-to-Video
For models that support video conditioning, upload the reference video with
`input_reference`:
```bash
curl -s http://localhost:8091/v1/videos \
-F "prompt=continue this motion with consistent subjects and lighting" \
-F "input_reference=@input.mp4;type=video/mp4" \
-F "width=1280" \
-F "height=720" \
-F "num_frames=80" \
-F "fps=16"
```
You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload
through `video_reference`. Do not send `video_reference` together with
`input_reference` or `image_reference`.
```bash
curl -s http://localhost:8091/v1/videos \
-F "prompt=continue this motion with consistent subjects and lighting" \
-F 'video_reference={"video_url":"https://example.com/input.mp4"}' \
-F "width=1280" \
-F "height=720" \
-F "num_frames=80" \
-F "fps=16"
```
JSON references currently support `image_url`/`video_url`; `file_id` references
are not implemented yet. Models may expose additional V2V controls through
`extra_params`. For example, Cosmos3 supports
`condition_frame_indexes_vision` and `condition_video_keep` to select which
decoded reference frames are used as clean conditioning.
### Synchronous Generation
```bash
@@ -184,10 +146,7 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
## Model-Specific Examples
For complete text-to-video, image-to-video, and model-specific video-to-video
walkthroughs, see:
For complete text-to-video and image-to-video walkthroughs, see:
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
for model-specific video-to-video examples and conditioning controls
-98
View File
@@ -1,98 +0,0 @@
# Production Metrics
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
```bash
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
curl http://localhost:8000/metrics
```
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
## Metric Namespaces
| Prefix | Source | Present when |
|--------|--------|--------------|
| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
| `http_` / `process_` | Uvicorn / Python runtime | Always |
## Pipeline-Level Metrics (`vllm:omni_`)
Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
### Request counts
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
### Latency
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
## Audio Modality Metrics (`vllm:omni_`)
Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description |
|--------|------|-------------|-------------|
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
## Cross-Stage Transfer Metrics (`vllm:omni_`)
Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
## vLLM Engine Metrics (`vllm:`)
When the pipeline includes an LLM stage, the upstream vLLM engine exposes its full set of ~37 metric families under the `vllm:` prefix.
vLLM-Omni wraps the upstream `vllm.v1.metrics.loggers.PrometheusStatLogger` with `OmniPrometheusStatLogger` so that the original `engine` single label is reshaped into `stage` + `replica`. Every `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — therefore gains per-`(stage, replica)` visibility automatically. No omni-side duplicate is needed for the text path.
```text
# Before wrap:
vllm:num_requests_running{model_name="...", engine="1"} 3.0
# After wrap:
vllm:num_requests_running{model_name="...", stage="1", replica="0"} 2.0
vllm:num_requests_running{model_name="...", stage="1", replica="1"} 1.0
```
For the full list of upstream metrics, see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md).
## Metric Availability by Pipeline Type
| Metric group | Multi-stage LLM (Qwen3-Omni) |
|---|---|
| `vllm:omni_` request tracking + latency | With `--log-stats` |
| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
## Naming Convention
- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix.
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
+1 -6
View File
@@ -133,27 +133,22 @@ The following tables show which models support each feature:
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
> Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
### VideoGen
| 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.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
**Frame Interpolation Support**
@@ -17,9 +17,6 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's
Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE
backend.
Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
To explicitly utilize a custom deployment YAML, specify the configuration path:
@@ -75,12 +72,6 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
```
To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server:
```bash
VLLM_USE_FLASHINFER_MOE_FP16=1 \
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
For the stage-based CLI, you usually do **not** need `--stage-overrides` for
that kind of change. Since each command launches one stage, just pass the knob
directly on that stage command:
@@ -380,31 +380,6 @@ curl -X POST http://localhost:8091/v1/audio/voices \
```
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
```bash
python qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio" \
--mode icl \
--output-dir /path/to/custom_voices
```
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
Then start the server with that config and call the Speech API with only the voice name:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
--output alice.wav
```
### Streaming PCM
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
@@ -444,10 +419,8 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
- Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers.
---
@@ -475,23 +448,6 @@ python voxcpm2/openai_speech_client.py \
```
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
### Precomputed custom voices
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
```bash
python voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--mode ref_continuation \
--prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices
```
Add the output directory to the deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
### Gradio demo (gapless streaming via AudioWorklet)
```bash
python voxcpm2/gradio_demo.py
-44
View File
@@ -28,50 +28,6 @@ in deep DiT blocks.
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available.
### Faster FP8 GEMM on Blackwell (quack)
On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which
applies the bias as a separate kernel after the GEMM. On the small GEMMs in video
DiTs this bias add is a significant overhead. Installing the optional `quack` kernel
lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering
that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster).
```bash
# CUDA 12.9
pip install vllm-omni[quack]
# CUDA 13.x
pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
```
It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
not installed, FP8 still works — it just keeps the unfused FlashInfer path.
#### Compile cache and warmup
quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer
the first time across all autotuned configs). The compiled `.o` files are cached on
disk and reused on later runs, so this is a one-time cost — **not** per request.
vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
the image — so the first cold start does not recompile. The engine's startup dummy
run already exercises the kernels, so with a warm cache the first real request is fast.
To pre-warm specific shapes (e.g. at image build time):
```python
from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8
# (M, K, N) per linear; M = number of tokens for your resolution/frame count
warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
```
> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install
> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12.
## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2)
+1 -1
View File
@@ -74,7 +74,7 @@ Online:
vllm serve Qwen/Qwen-Image \
--omni \
--port 8000 \
--diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
--quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
```
## Parameters
@@ -195,101 +195,6 @@ stages:
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.
+6 -2
View File
@@ -121,6 +121,9 @@ def parse_args():
help="Temperature for text generation sampling (default: 0.3).",
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
args = parser.parse_args()
return args
@@ -151,7 +154,6 @@ def main():
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"
@@ -162,7 +164,9 @@ def main():
if args.quantization:
omni_kwargs["quantization_config"] = args.quantization
omni = Omni(model=model_name, **omni_kwargs)
# Override CLI --model with the derived model_name.
omni_kwargs["model"] = model_name
omni = Omni(**omni_kwargs)
formatted_prompts = []
for p in prompts:
@@ -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"
}
@@ -1,304 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import argparse
import uuid
from pathlib import Path
import cv2
import numpy as np
import torch
from PIL import Image
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
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:
repo_root = Path(__file__).resolve().parents[3]
parser = argparse.ArgumentParser(description="Export DreamZero prediction video from downloaded example inputs.")
parser.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
parser.add_argument("--deploy-config", type=Path, required=True)
parser.add_argument(
"--video-dir",
type=Path,
default=repo_root / "outputs" / "dreamzero" / "assets",
help="Directory containing the three camera MP4 files.",
)
parser.add_argument(
"--output-dir", type=Path, default=repo_root / "outputs" / "dreamzero" / "generated_predictions"
)
parser.add_argument("--output-stem", default="dreamzero_prediction")
parser.add_argument(
"--prompt",
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("--save-input-video", action="store_true")
parser.add_argument("--save-gif", action="store_true")
parser.add_argument("--save-actions", action="store_true")
parser.add_argument("--fps", type=int, default=5)
return parser.parse_args()
def _load_all_frames(video_path: Path) -> np.ndarray:
cap = cv2.VideoCapture(str(video_path))
frames = []
while True:
ok, frame = cap.read()
if not ok:
break
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
cap.release()
if not frames:
raise RuntimeError(f"No frames loaded from {video_path}")
return np.stack(frames, axis=0)
def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
camera_frames: dict[str, np.ndarray] = {}
for camera_key, file_name in CAMERA_FILES.items():
video_path = video_dir / file_name
if not video_path.exists():
raise FileNotFoundError(
f"Missing DreamZero example asset: {video_path}. "
"Download the example videos with: "
f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
)
camera_frames[camera_key] = _load_all_frames(video_path)
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],
*,
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]]:
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,
)
)
if len(observations) < 2:
raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
return camera_frames, observations[:2]
def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
if not isinstance(output, OmniRequestOutput):
raise TypeError(f"Expected OmniRequestOutput, got {type(output)!r}")
if not output.images:
raise RuntimeError("DreamZero output does not contain video latents in `images`.")
latents = output.images[0]
if not isinstance(latents, torch.Tensor):
raise TypeError(f"Expected tensor latents, got {type(latents)!r}")
latents = latents.detach().cpu()
if latents.dim() == 4:
latents = latents.unsqueeze(0)
if latents.dim() != 5:
raise ValueError(f"Unexpected latent shape: {tuple(latents.shape)}")
if latents.shape[1] < latents.shape[2]:
latents = latents.transpose(1, 2).contiguous()
return latents
def _write_mp4(path: Path, frames: np.ndarray, fps: int) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
height, width = frames.shape[1:3]
writer = cv2.VideoWriter(
str(path),
cv2.VideoWriter_fourcc(*"mp4v"),
float(fps),
(width, height),
)
if not writer.isOpened():
raise RuntimeError(f"Failed to open video writer for {path}")
try:
for frame in frames:
writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
finally:
writer.release()
def _write_gif(path: Path, frames: np.ndarray, fps: int) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
images = [Image.fromarray(frame) for frame in frames]
duration_ms = max(int(round(1000 / max(fps, 1))), 1)
images[0].save(
path,
save_all=True,
append_images=images[1:],
duration=duration_ms,
loop=0,
)
def _stitch_input_frames(camera_frames: dict[str, np.ndarray]) -> np.ndarray:
total_frames = min(frames.shape[0] for frames in camera_frames.values())
stitched = []
for frame_index in range(total_frames):
left = camera_frames["observation/exterior_image_0_left"][frame_index]
right = camera_frames["observation/exterior_image_1_left"][frame_index]
wrist = camera_frames["observation/wrist_image_left"][frame_index]
pad = np.zeros((left.shape[0], left.shape[1], 3), dtype=np.uint8)
canvas = np.concatenate([left, right], axis=1)
bottom = np.concatenate([wrist, pad], axis=1)
stitched.append(np.concatenate([canvas, bottom], axis=0))
return np.stack(stitched, axis=0)
def _run_generation(
model: str, deploy_config_path: Path, observations: list[dict]
) -> tuple[Omni, list[OmniRequestOutput]]:
omni = Omni(
model=model,
deploy_config=str(deploy_config_path),
enforce_eager=True,
worker_extension_cls=WORKER_EXTENSION,
)
outputs: list[OmniRequestOutput] = []
for index, obs in enumerate(observations):
sampling_params = OmniDiffusionSamplingParams(
extra_args={
"reset": index == 0,
"session_id": obs["session_id"],
"robot_obs": obs,
}
)
result = omni.generate(obs["prompt"], sampling_params_list=[sampling_params])
if not result:
raise RuntimeError(f"No output returned for DreamZero request {index}")
outputs.append(result[0])
return omni, outputs
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)
if engine is None:
raise RuntimeError("DreamZero export requires inline diffusion stage access.")
decoded = engine.executor.collective_rpc(
"decode_video_latents_to_uint8",
args=(full_latents,),
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 main() -> None:
args = _parse_args()
session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
camera_frames, observations = _build_observations(
video_dir=args.video_dir,
prompt=args.prompt,
session_id=session_id,
)
args.output_dir.mkdir(parents=True, exist_ok=True)
if args.save_input_video:
input_frames = _stitch_input_frames(camera_frames)
_write_mp4(args.output_dir / f"{args.output_stem}_input.mp4", input_frames, fps=15)
if args.save_gif:
_write_gif(args.output_dir / f"{args.output_stem}_input.gif", input_frames[::3], fps=5)
omni = None
try:
omni, outputs = _run_generation(
model=args.model,
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)
finally:
if omni is not None:
omni.close()
mp4_path = args.output_dir / f"{args.output_stem}.mp4"
_write_mp4(mp4_path, frames, fps=args.fps)
print(f"SAVED_MP4={mp4_path}")
if args.save_gif:
gif_path = args.output_dir / f"{args.output_stem}.gif"
_write_gif(gif_path, frames, fps=args.fps)
print(f"SAVED_GIF={gif_path}")
if args.save_actions:
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")),
)
print(f"SAVED_ACTIONS={npz_path}")
if __name__ == "__main__":
main()
@@ -42,7 +42,7 @@ pip install \
## 3. Hardware and VRAM Requirements
This example uses a 3-stage pipeline on one GPU by default
([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)):
([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
@@ -18,6 +18,8 @@ import numpy as np
import torch
from PIL import Image
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
TASK_DEFAULT_RUNTIME = {
@@ -878,8 +880,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument(
"--stage-config-path",
type=str,
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
help="Path to deploy config yaml.",
default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"),
help="Path to stage config yaml.",
)
parser.add_argument(
"--dynin-config-path",
@@ -970,6 +972,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -107,6 +107,9 @@ def parse_args():
),
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -143,20 +143,6 @@ def parse_args() -> argparse.Namespace:
default=None,
required=False,
)
parser.add_argument(
"--width",
type=int,
default=None,
metavar="W",
help="Output image width in pixels. Default: None (pipeline's default).",
)
parser.add_argument(
"--height",
type=int,
default=None,
metavar="H",
help="Output image height in pixels. Default: None (pipeline's default).",
)
parser.add_argument(
"--seed",
type=int,
@@ -250,8 +236,8 @@ def parse_args() -> argparse.Namespace:
parser.add_argument(
"--resolution",
type=int,
default=None,
help="Bucket in (640, 1024) to determine the condition and output resolution. If width and height are not provided, this will be set to default 640.",
default=640,
help="Bucket in (640, 1024) to determine the condition and output resolution",
)
parser.add_argument(
@@ -398,15 +384,6 @@ def parse_args() -> argparse.Namespace:
def main():
args = parse_args()
if args.resolution and (args.width or args.height):
raise ValueError("--resolution and --width/--height cannot be specified together")
if args.width is not None and args.width <= 0:
raise ValueError("--width must be a positive integer")
if args.height is not None and args.height <= 0:
raise ValueError("--height must be a positive integer")
if not args.width and not args.height and not args.resolution:
args.resolution = 640
# Validate input images exist and load them
input_images = []
for image_path in args.image:
@@ -487,7 +464,6 @@ def main():
print(
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
)
print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}")
print(f"{'=' * 60}\n")
generation_start = time.perf_counter()
@@ -497,26 +473,22 @@ def main():
omni.start_profile()
# Generate edited image
sampling_params = OmniDiffusionSamplingParams(
width=args.width,
height=args.height,
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_2,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_outputs_per_prompt,
layers=args.layers,
)
if args.resolution:
sampling_params.resolution = args.resolution
outputs = omni.generate(
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
"multi_modal_data": {"image": input_image},
},
sampling_params,
OmniDiffusionSamplingParams(
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_2,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_outputs_per_prompt,
layers=args.layers,
resolution=args.resolution,
),
)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
@@ -210,12 +210,6 @@ def main():
args = parse_args()
os.makedirs(args.output, exist_ok=True)
# Lance aligns more naturally with upstream when using FLASH_ATTN
# (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
# numerical drift through the 36-layer Qwen2 stack on B300). Default
# to flash_attn when the user hasn't explicitly set a backend.
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
if args.txt_prompts:
with open(args.txt_prompts, encoding="utf-8") as f:
prompts = [ln.strip() for ln in f if ln.strip()]
@@ -288,19 +288,11 @@ def _init_models(
f" replicas={replicas_per_omni}...",
flush=True,
)
# Lance single-stage engine defaults — mirrors the offline ``end2end.py``
# flat-kwargs path so this demo does not depend on ``deploy/lance.yaml``.
# ``pipeline="lance"`` is the explicit selector (Lance HF config has no
# ``model_type`` field, so auto-detect from the model dir alone fails).
lance_engine_defaults = dict(
pipeline="lance",
max_num_batched_tokens=32768,
max_num_seqs=1,
enforce_eager=True,
trust_remote_code=True,
enable_prefix_caching=False,
async_chunk=False,
)
# Lance needs ``vllm_omni/deploy/lance.yaml`` as the pipeline-selector
# (Lance HF config has no ``model_type`` field for auto-detect). Pass
# via ``deploy_config`` — flat kwargs alone don't trigger pipeline
# resolution.
lance_engine_defaults = dict(deploy_config="vllm_omni/deploy/lance.yaml")
img_kwargs: dict = {"model": img_ckpt, **lance_engine_defaults}
if img_devices != "0":
@@ -1245,11 +1237,6 @@ def build_ui(
def main():
args = parse_args()
# Lance aligns better with upstream when using FLASH_ATTN
# (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
# numerical drift through the 36-layer Qwen2 stack). Default to
# flash_attn unless the user explicitly set a backend.
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
print("=" * 60, flush=True)
print("[unified] Lance × vllm-omni demo", flush=True)
print(f" snapshot: {args.model}", flush=True)
@@ -19,6 +19,7 @@ from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
@@ -48,6 +49,7 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -29,6 +29,7 @@ from PIL import Image
from vllm.sampling_params import SamplingParams
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
@@ -117,6 +118,7 @@ def parse_args() -> argparse.Namespace:
)
p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
nullify_stage_engine_defaults(p)
args = p.parse_args()
if not args.prompt:
args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
@@ -21,10 +21,11 @@ from message_convert import (
to_prompt,
)
from vllm import SamplingParams
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.inputs.data import OmniTokensPrompt
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SEED = 42
@@ -342,7 +343,7 @@ def main(args):
def parse_args():
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--model-name",
"-m",
@@ -438,6 +439,7 @@ def parse_args():
"vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -53,7 +53,7 @@ python examples/offline_inference/ming_flash_omni/end2end.py \
--image-path ./3_0.png
```
### Omni-Speech (Thinker + Talker) {#omni-speech-thinker--talker}
### Omni-Speech (Thinker + Talker)
The default deploy YAML already runs thinker+talker, so spoken output only requires requesting `audio` (or `text,audio`) modalities.
The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.
@@ -17,13 +17,14 @@ from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.media.audio import load_audio
from vllm.utils.argparse_utils import FlexibleArgumentParser
import vllm_omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
# Imports the processor also registers itself
from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor # noqa: F401
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
SEED = 42
MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
@@ -396,7 +397,7 @@ def main(args):
def parse_args():
parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
parser.add_argument(
"--query-type",
"-q",
@@ -495,6 +496,7 @@ def parse_args():
help="Output directory for results.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -19,9 +19,10 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.media.audio import load_audio
from vllm.sampling_params import SamplingParams
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.utils.tracking_parser import TrackingArgumentParser
SEED = 42
@@ -431,7 +432,7 @@ def main(args):
def parse_args():
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--model",
type=str,
@@ -553,6 +554,7 @@ def parse_args():
default=False,
help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -19,9 +19,10 @@ from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
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.utils.tracking_parser import TrackingArgumentParser
SEED = 42
@@ -421,7 +422,7 @@ def main(args):
def parse_args():
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--model",
type=str,
@@ -560,6 +561,7 @@ def parse_args():
help="Model dtype (auto, half, float16, bfloat16, float, float32).",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -39,9 +39,10 @@ from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
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.async_omni import AsyncOmni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
@@ -382,9 +383,7 @@ async def run_all(args):
print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
async_omni = None
try:
# ``from_cli_args`` forwards only explicitly-passed CLI args so
# argparse defaults do not silently override deploy YAML values.
async_omni = AsyncOmni.from_cli_args(args, model=args.model)
async_omni = AsyncOmni(**vars(args))
# Use default sampling params from stage config (they are pre-configured
# in the YAML for each stage).
@@ -454,7 +453,7 @@ async def run_all(args):
def parse_args():
parser = TrackingArgumentParser(
parser = FlexibleArgumentParser(
description=(
"Offline inference with async_chunk enabled via AsyncOmni. "
"Downstream stages start before upstream stages finish, "
@@ -590,6 +589,7 @@ def parse_args():
default=16000,
help="Sampling rate for audio loading.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -29,6 +29,7 @@ import os
from PIL import Image
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -167,6 +168,7 @@ def parse_args():
help="Enable cache-dit summary logging after diffusion forward passes.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -16,10 +16,9 @@ Usage:
"""
import argparse
import json
import os
import time
from pathlib import Path
from typing import Any
import torch
@@ -30,16 +29,6 @@ from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
return config
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Generate a talking-head video from a reference image and audio (Wan2.2 S2V)."
@@ -98,12 +87,6 @@ def parse_args() -> argparse.Namespace:
default=3.0,
help="Scheduler flow shift (default: 3.0).",
)
parser.add_argument(
"--boundary-ratio",
type=float,
default=None,
help="Boundary split ratio for low/high DiT stages (S2V). Default varies by model.",
)
parser.add_argument(
"--output",
type=str,
@@ -195,10 +178,30 @@ def parse_args() -> argparse.Namespace:
help="Enable cache-dit summary logging after diffusion forward passes.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
"--profile-dir",
type=str,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
help="Enable torch profiler and save traces to this directory.",
)
parser.add_argument(
"--profile-record-shapes",
action="store_true",
help="Record tensor shapes in profiler (increases trace size).",
)
parser.add_argument(
"--profile-with-stack",
action="store_true",
help="Record stack traces in profiler (increases overhead).",
)
parser.add_argument(
"--profile-with-memory",
action="store_true",
help="Profile memory usage.",
)
parser.add_argument(
"--profile-with-flops",
action="store_true",
help="Estimate FLOPs for operations.",
)
return parser.parse_args()
@@ -236,9 +239,23 @@ def main():
ulysses_degree=args.ulysses_degree,
)
profiler_enabled = args.profiler_config is not None
# Check if profiling is requested via CLI args or environment variable
profile_dir = args.profile_dir or os.getenv("VLLM_TORCH_PROFILER_DIR")
profiler_enabled = bool(profile_dir)
profiler_config = None
if profiler_enabled:
from vllm.config import ProfilerConfig
omni_kwargs = dict(
profiler_config = ProfilerConfig(
profiler="torch",
torch_profiler_dir=profile_dir,
torch_profiler_record_shapes=args.profile_record_shapes,
torch_profiler_with_stack=args.profile_with_stack,
torch_profiler_with_memory=args.profile_with_memory,
torch_profiler_with_flops=args.profile_with_flops,
)
omni = Omni(
model=args.model,
model_class_name="WanS2VPipeline",
flow_shift=args.flow_shift,
@@ -252,12 +269,8 @@ def main():
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_cache_dit_summary=args.enable_cache_dit_summary,
profiler_config=args.profiler_config,
profiler_config=profiler_config,
)
if args.boundary_ratio is not None:
omni_kwargs["boundary_ratio"] = args.boundary_ratio
omni = Omni(**omni_kwargs)
# Print generation configuration
print(f"\n{'=' * 60}")
@@ -312,16 +325,13 @@ def main():
if profiler_enabled:
print("\n[Profiler] Stopping profiler and collecting results...")
profile_results = omni.stop_profile()
if profile_results and isinstance(profile_results, dict):
traces = profile_results.get("traces", [])
if profile_results and isinstance(profile_results, list):
print("\n" + "=" * 60)
print("PROFILING RESULTS:")
for rank, trace in enumerate(traces):
print(f"\nRank {rank}:")
if trace:
print(f" • Trace: {trace}")
if not traces:
print(" No traces collected.")
for i, result in enumerate(profile_results):
print(f"\nStage {i}:")
if result:
print(f" {result}")
print("=" * 60)
else:
print("[Profiler] No valid profiling data returned.")
@@ -307,6 +307,9 @@ def parse_args() -> argparse.Namespace:
help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
)
current_omni_platform.pre_register_and_update(parser)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -17,7 +17,6 @@ list of supported architectures across all modalities, see
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
@@ -160,46 +159,6 @@ Streaming requires `async_chunk: true` in the stage config.
---
## Ming-omni-tts
Dense 0.5B two-stage TTS pipeline (`AR + flow` + audio VAE) at 44.1 kHz. The example covers style, IP voice, music-only generation, text-to-audio events, emotion, dialect, zero-shot cloning, podcast, speech+BGM, and speech+environment-sound cases.
### Quick start
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case style \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
### Voice cloning
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case zero_shot \
--ref-audio /path/to/reference.wav \
--ref-text "在此奉劝大家别乱打美白针。" \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
### Streaming
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case basic \
--ref-audio /path/to/reference.wav \
--streaming \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
### Notes
- `style`, `ip`, `bgm`, and `tta` do not require reference audio.
- Reference-audio cases use `--ref-audio`; `zero_shot` also requires `--ref-text`.
- `podcast` uses multiple references via `--ref-audio-paths`.
- Full case details live in [`ming_tts/README.md`](ming_tts/README.md).
---
## Ming-flash-omni-TTS
Standalone talker-only deployment of Ming-flash-omni-2.0 at 44.1 kHz. Voice is controlled through caption fields (`风格` / `IP` / `语速`/`基频`/`音量`) rather than reference audio.
@@ -10,6 +10,7 @@ import soundfile as sf
from vllm import SamplingParams
from vllm.multimodal.media.audio import load_audio
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
@@ -64,6 +65,7 @@ def run_e2e():
required=True,
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
)
nullify_stage_engine_defaults(parser)
args = parser.parse_args()
# Ensure tokenizer directory exists
if not os.path.exists(args.tokenizer):
@@ -26,14 +26,16 @@ import torch
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
build_fish_text_only_prompt_ids,
estimate_fish_voice_clone_prompt_len_from_normalized,
normalize_fish_voice_clone_texts,
)
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
@@ -209,7 +211,7 @@ async def main_streaming(args):
def parse_args():
parser = TrackingArgumentParser(
parser = FlexibleArgumentParser(
description="Fish Speech S2 Pro offline inference with vLLM Omni",
)
parser.add_argument(
@@ -265,6 +267,7 @@ def parse_args():
default=False,
help="Stream audio chunks as they arrive via AsyncOmni.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -30,9 +30,10 @@ from pathlib import Path
import numpy as np
import soundfile as sf
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
SAMPLE_RATE = 24_000
DEFAULT_TEXTS = (
@@ -42,7 +43,7 @@ DEFAULT_TEXTS = (
def parse_args():
parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference")
parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
parser.add_argument(
"--model",
type=str,
@@ -87,6 +88,7 @@ def parse_args():
default=None,
help="Transcript of the reference clip. Required when --ref-audio is set.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -1,204 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline higgs-audio v3 inference example.
Runs Stage 0 (Qwen3 talker) + Stage 1 (HiggsAudio codec) end-to-end
through the vLLM-Omni engine without going through the HTTP server, and
saves a 24 kHz mono WAV per prompt.
Example (plain TTS):
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
--texts "Hello world." \\
--output-dir results/higgs_v3_wavs
Example (voice clone):
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
--texts "Hello world." \\
--ref-audio path/to/reference.wav \\
--ref-text "Transcript of the reference clip." \\
--output-dir results/higgs_v3_wavs
"""
from __future__ import annotations
import os
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
import re
import time
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
from vllm_omni import Omni
SAMPLE_RATE = 24_000
DEFAULT_TEXTS = (
"Hello world.",
"The quick brown fox jumps over the lazy dog.",
"Today is a beautiful day for a walk in the park.",
)
def parse_args():
import argparse
parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
parser.add_argument(
"--model",
type=str,
default="bosonai/higgs-audio-v3-tts-4b",
help="Stage-0 talker model id or path.",
)
parser.add_argument(
"--texts",
type=str,
nargs="+",
default=list(DEFAULT_TEXTS),
help="One or more plain-text prompts to synthesize.",
)
parser.add_argument(
"--output-dir",
type=str,
default="results/higgs_v3_wavs",
help="Directory to write per-prompt WAV files.",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override the deploy config path.",
)
parser.add_argument(
"--max-new-tokens",
type=int,
default=2048,
help="Cap on Stage-0 codec frames per request.",
)
parser.add_argument(
"--ref-audio",
type=str,
default=None,
help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Transcript of the reference audio. Optional but improves fidelity.",
)
return parser.parse_args()
def _slugify(text: str) -> str:
slug = re.sub(r"\s+", "_", text.strip().lower())
slug = re.sub(r"[^a-z0-9_]+", "", slug)
return slug[:48] or "prompt"
def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
audio = multimodal_output.get("model_outputs")
if audio is None:
audio = multimodal_output.get("audio")
if audio is None:
raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
if isinstance(audio, list):
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
if not valid:
raise ValueError("audio list is empty")
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
return torch.as_tensor(audio).float().cpu().reshape(-1)
def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
arr = pcm.numpy()
if arr.dtype.kind == "f":
arr = np.clip(arr, -1.0, 1.0)
arr = (arr * 32767.0).astype(np.int16)
else:
arr = arr.astype(np.int16)
return arr
def main():
args = parse_args()
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
apply_delay_pattern,
encode_reference_audio,
)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
# Load and encode reference audio once (if voice cloning)
ref_codes_delayed: torch.Tensor | None = None
if args.ref_audio is not None:
ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
if ref_wav.ndim == 2:
ref_wav = ref_wav.mean(axis=1)
ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
print(f"Reference : {args.ref_audio}")
print(f"Ref codes : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
if args.ref_text:
print(f"Ref text : {args.ref_text}")
print(f"Model : {args.model}")
print(f"Prompts : {len(args.texts)}")
print(f"Output dir : {output_dir}")
print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")
total_elapsed = 0.0
total_dur = 0.0
for text in args.texts:
if ref_codes_delayed is not None:
prompt_ids = adapter.build_prompt(
text,
num_ref_tokens=int(ref_codes_delayed.shape[0]),
reference_text=args.ref_text,
)
prompt = {
"prompt_token_ids": prompt_ids,
"additional_information": {
"audio_input_ids": ref_codes_delayed.to(torch.long),
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
},
}
else:
prompt_ids = adapter.build_prompt(text)
prompt = {"prompt_token_ids": prompt_ids}
t_start = time.perf_counter()
outputs = engine.generate([prompt])
elapsed = time.perf_counter() - t_start
total_elapsed += elapsed
mm = outputs[0].outputs[0].multimodal_output
pcm = _extract_pcm(mm)
slug = _slugify(text)
out_path = output_dir / f"{slug}.wav"
sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
dur = pcm.numel() / SAMPLE_RATE
total_dur += dur
print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}")
rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}")
if __name__ == "__main__":
main()
@@ -5,11 +5,11 @@ from typing import Any
import soundfile as sf
import torch
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
from vllm.utils.argparse_utils import FlexibleArgumentParser
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
@@ -68,7 +68,7 @@ def save_audio(mm: dict[str, Any], output_path: str) -> None:
def parse_args():
parser = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
parser.add_argument(
"--deploy-config",
@@ -89,6 +89,7 @@ def parse_args():
parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -1,248 +0,0 @@
# Ming-omni-tts Offline Inference
`end2end.py` runs Ming dense 0.5B end to end with vLLM-Omni. It uses the in-repo Ming prompt assembly helper directly, so the example request shape matches the real integration instead of a simplified wrapper.
## Files
| File | Purpose |
|---|---|
| `end2end.py` | Driver: CLI, case loading, prompt construction, orchestration (~150 lines) |
| `cases.yaml` | All 11 built-in case definitions (prompt, text, instruction, ref-audio flags, flow controls) |
| `runner.py` | Engine management and audio output (streaming + blocking paths) |
## Model Overview
Ming dense 0.5B is exposed here as a two-stage offline pipeline:
- **Stage 0**: Qwen2-based AR generation with Ming prompt formatting and inline flow controls
- **Stage 1**: audio VAE decode to mono 44.1 kHz waveform
`config_ming_tts.py` adapts the checkpoint's HuggingFace config fields
(LLM, DiT, aggregator, AudioVAE, token ids). `vllm_omni/deploy/ming_tts.yaml`
selects the vLLM-Omni pipeline and stage runtime topology, including
connectors, async chunking, memory limits, and sampling defaults.
The example supports both:
- **Blocking eager** via `vllm_omni/deploy/ming_tts.yaml`
- **Async chunk eager** via `vllm_omni/deploy/ming_tts.yaml` (default `async_chunk: true`)
## Setup
Install vLLM-Omni with the platform requirements for your accelerator:
```bash
uv pip install -e .
```
The Ming offline example does not require a separate upstream Ming package.
Reference-audio cases use the repo dependencies for audio loading,
resampling, and CampPlus speaker extraction, including `soundfile`,
`torchaudio`, and `onnxruntime-rocm` in the documented ROCm environment.
The tested ROCm environment is summarized in the
[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
## Supported Cases
These cases cover the upstream dense 0.5B cookbook surface that maps cleanly onto the current vLLM-Omni example:
- `style`: zero-speaker style-conditioned speech
- `ip`: zero-speaker IP voice generation
- `bgm`: music-only generation
- `tta`: text-to-audio event generation with FlowLoss controls
- `emotion`: reference-audio speech with emotion control
- `basic`: reference-audio speech with speed / pitch / volume control
- `dialect`: reference-audio speech with dialect control
- `zero_shot`: reference-audio cloning with explicit transcript
- `podcast`: multi-reference dialogue generation with automatic speaker embedding extraction
- `speech_bgm`: speech with background music conditioning
- `speech_sound`: speech with environmental sound conditioning
## Quick Start
Run the zero-speaker style example:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case style \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
Run zero-shot cloning with a transcript:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case zero_shot \
--ref-audio /path/to/10002287-00000094.wav \
--ref-text "在此奉劝大家别乱打美白针。" \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
Run emotion-controlled speech:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case emotion \
--ref-audio /path/to/emotion_prompt.wav \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
Run podcast generation with two reference clips:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case podcast \
--ref-audio-paths /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
The script automatically extracts one 192-d speaker embedding per reference WAV using the Ming model's `campplus.onnx`.
If you already have precomputed multi-speaker embeddings, you can override extraction with:
```bash
--speaker-embedding /path/to/podcast_speaker_embeddings.json
```
where the JSON is a list of speaker embeddings, one 192-d vector per speaker.
Run text-to-audio event generation:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case tta \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
Use async_chunk streaming:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case basic \
--ref-audio /path/to/10002287-00000095.wav \
--streaming \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager
```
`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently
supports one prompt per process invocation; use blocking mode for
`--num-prompts > 1`.
Collect runtime stats and a manifest:
```bash
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--case style \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--enforce-eager \
--enable-stats \
--stats-log-file output_audio/ming_style_pipeline.log \
--metadata-json output_audio/ming_style_manifest.json
```
## Reference Fixtures
The upstream Ming cookbook uses these public audio fixtures from `inclusionAI/Ming-omni-tts/data/wavs`:
- `10002287-00000094.wav` for zero-shot cloning
- `10002287-00000095.wav` for `basic`
- `emotion_prompt.wav` for `emotion`
- `yue_prompt.wav` for `dialect`
- `00000309-00000300.wav` for `speech_bgm` and `speech_sound`
- `CTS-CN-F2F-2019-11-11-423-012-A.wav` and `CTS-CN-F2F-2019-11-11-423-012-B.wav` for `podcast`
## Validation Matrix
The repo-facing example is intended to cover the same dense TTS workflows used
by the local Ming validation script:
| Case | Blocking `deploy/ming_tts.yaml` | Async chunk `deploy/ming_tts.yaml` | Extra inputs |
|---|---:|---:|---|
| `style` | Yes | Optional smoke test | none |
| `ip` | Yes | Optional smoke test | none |
| `bgm` | Yes | Optional smoke test | none |
| `tta` | Yes | Optional smoke test | none |
| `emotion` | Yes | Yes | `--ref-audio emotion_prompt.wav` |
| `basic` | Yes | Yes | `--ref-audio 10002287-00000095.wav` |
| `dialect` | Yes | Yes | `--ref-audio yue_prompt.wav` |
| `zero_shot` | Yes | Yes | `--ref-audio 10002287-00000094.wav --ref-text ...` |
| `podcast` | Yes | Yes | two `--ref-audio-paths` |
| `speech_bgm` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
| `speech_sound` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
## Validated Outputs
The following measurements are retained from an earlier L4 CUDA validation;
they are not ROCm benchmark results. Default async_chunk matched blocking
output frame counts and Stage-1 patch counts for every case:
| Case | Blocking frames / patches / sec | Async chunk frames / patches / sec |
|---|---:|---:|
| `style` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
| `ip` | 183456 / 13 / 4.16 | 183456 / 13 / 4.16 |
| `bgm` | 1326528 / 94 / 30.08 | 1326528 / 94 / 30.08 |
| `tta` | 465696 / 33 / 10.56 | 465696 / 33 / 10.56 |
| `emotion` | 324576 / 23 / 7.36 | 324576 / 23 / 7.36 |
| `basic` | 211680 / 15 / 4.80 | 211680 / 15 / 4.80 |
| `dialect` | 239904 / 17 / 5.44 | 239904 / 17 / 5.44 |
| `zero_shot` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
| `podcast` | 437472 / 31 / 9.92 | 437472 / 31 / 9.92 |
| `speech_bgm` | 296352 / 21 / 6.72 | 296352 / 21 / 6.72 |
| `speech_sound` | 352800 / 25 / 8.00 | 352800 / 25 / 8.00 |
## Key Arguments
| Argument | Description |
|---|---|
| `--model` | Hugging Face repo or local Ming checkpoint path |
| `--deploy-config` | Deploy config YAML. Use `vllm_omni/deploy/ming_tts.yaml` |
| `--case` | Built-in demo case |
| `--ref-audio` | Single reference wav path for cloning-style cases |
| `--ref-audio-paths` | Multiple reference wav paths, used by `podcast` |
| `--ref-text` | Reference transcript. Required for `zero_shot` |
| `--instructions` | Free-form Ming instruction string |
| `--instruction-json` | Structured Ming instruction JSON |
| `--speaker-embedding` | JSON file containing a 192-d speaker embedding |
| `--extract-speaker-embeddings` | Force CampPlus speaker extraction from the provided reference audio paths |
| `--max-decode-steps` | Override `ming_max_decode_steps` |
| `--num-prompts` | Repeat the same case N times. Outputs are indexed when `N > 1` |
| `--streaming` | Use `AsyncOmni` and async_chunk transport |
| `--enforce-eager` | Recommended for Ming dense; non-eager is out of scope |
| `--enable-stats` / `--log-stats` | Enable vLLM-Omni per-request stats logging |
| `--stats-log-file` | Optional path for the stats log |
| `--metadata-json` | Optional path for the run manifest JSON |
| `--stage-init-timeout` | Per-stage initialization timeout in seconds |
| `--init-timeout` | Total initialization timeout in seconds |
| `--batch-timeout` | Batch timeout in seconds |
| `--worker-backend` | `multi_process` or `ray` |
| `--ray-address` | Ray cluster address when using `--worker-backend ray` |
## Output
- The script writes one mono 44.1 kHz WAV file per run
- Default output directory: `output_audio/`
- Default filename: `ming_<case>.wav`
- When `--num-prompts > 1`, outputs are indexed as `ming_<case>_00000.wav`, `..._00001.wav`, etc.
- When stats are enabled, the script can also write:
- a stats log file such as `ming_style_pipeline.log`
- a manifest JSON with per-output metadata, stage durations, peak memory info,
and streaming client latency metrics when `--streaming` is used
## Notes
- `style` and `ip` are zero-speaker paths and do not require a reference clip
- `emotion`, `basic`, `dialect`, `speech_bgm`, and `speech_sound` require one reference clip
- `zero_shot` requires both `--ref-audio` and `--ref-text`
- `podcast` requires at least two reference clips via `--ref-audio-paths`
- `podcast` automatically extracts one speaker embedding per reference clip
- `--speaker-embedding` may contain either one 192-d vector or a list of 192-d vectors
- `--enforce-eager` was used for the validated runs
- The earlier L4 validation used SDPA for the Ming audio VAE instead of
FlashAttention2, which is the preferred default when available.
@@ -1,110 +0,0 @@
style:
prompt: "Please generate speech based on the following description.\n"
text: "我会一直在这里陪着你,直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗?"
instruction:
风格: >-
这是一种ASMR耳语,属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语,声音气音成分重。音量极低,紧贴麦克风,语速极慢,旨在制造触发听者颅内快感的声学刺激。
use_zero_spk_emb: true
max_decode_steps: 200
ip:
prompt: "Please generate speech based on the following description.\n"
text: "这款产品的名字,叫变态坑爹牛肉丸。"
instruction:
IP: "灵小甄"
use_zero_spk_emb: true
max_decode_steps: 200
bgm:
prompt: "Please generate music based on the following description.\n"
text: " Genre: 电子舞曲. Mood: 自信 / 坚定. Instrument: 架子鼓. Theme: 节日. Duration: 30s."
instruction: null
use_zero_spk_emb: false
max_decode_steps: 400
tta:
prompt: "Please generate audio events based on given text.\n"
text: "Thunder and a gentle rain"
instruction: null
use_zero_spk_emb: false
max_decode_steps: 200
cfg: 4.5
sigma: 0.3
temperature: 2.5
emotion:
prompt: "Please generate speech based on the following description.\n"
text: "我竟然抢到了陈奕迅的演唱会门票!太棒了!终于可以现场听一听他的歌声了!"
instruction:
情感: "高兴"
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
basic:
prompt: "Please generate speech based on the following description.\n"
text: "简单地说,这相当于惠普把消费领域市场拱手相让了。"
instruction:
语速: "快速"
基频: "中"
音量: "中"
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
dialect:
prompt: "Please generate speech based on the following description.\n"
text: "我觉得社会企业同个人都有责任"
instruction:
方言: "广粤话"
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
zero_shot:
prompt: "Please generate speech based on the following description.\n"
text: "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。"
instruction: null
requires_ref_audio: true
requires_ref_text: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
podcast:
prompt: "Please generate speech based on the following description.\n"
text: " speaker_1:你可以说一下,就大概说一下,可能虽然我也不知道,我看过那部电影没有。\n speaker_2:就是那个叫什么,变相一节课的嘛。\n speaker_1:嗯。\n speaker_2:一部搞笑的电影。\n speaker_1:一部搞笑的。\n"
instruction: null
prompt_text: " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试,做个,当服务员还要去笔试了\n speaker_2:对啊,这真的很奇怪,就是 单纯的因,单纯自己工资不高,只是因为可能人家那个店比较出名一点,就对你苛刻要求\n"
requires_ref_audio_count: 2
auto_extract_speaker_embeddings: true
max_decode_steps: 200
speech_bgm:
prompt: "Please generate speech based on the following description.\n"
text: "此次业绩下滑原因,可归结为企业停止服务某些品牌,而带来的负面影响。"
instruction:
BGM:
Genre: "当代古典音乐."
Mood: "温暖 / 友善."
Instrument: "电吉他"
Theme: "节日."
SNR: 10.0
ENV: null
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
speech_sound:
prompt: "Please generate speech based on the following description.\n"
text: "此次业绩下滑原因,可归结为企业停止服务某些品牌,而带来的负面影响。"
instruction:
BGM:
ENV: "Birds chirping"
SNR: 10.0
Genre: null
Mood: null
Instrument: null
Theme: null
requires_ref_audio: true
auto_extract_speaker_embeddings: true
max_decode_steps: 200
@@ -1,200 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline inference demo for Ming-omni-tts via vLLM Omni."""
import json
import os
from pathlib import Path
import soundfile as sf
import torch
import torchaudio
import yaml
from transformers import AutoTokenizer
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
KEY_CFG,
KEY_MAX_DECODE_STEPS,
KEY_SIGMA,
KEY_TEMPERATURE,
SAMPLE_RATE,
)
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
try:
from .runner import (
build_manifest,
build_sampling_params,
resolve_metadata_json,
resolve_stats_log_file,
run_generation,
)
except ImportError:
from runner import (
build_manifest,
build_sampling_params,
resolve_metadata_json,
resolve_stats_log_file,
run_generation,
)
_DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
_REPO_ROOT = Path(__file__).resolve().parents[4]
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni/deploy/ming_tts.yaml")
_CASES_FILE = Path(__file__).with_name("cases.yaml")
CASE_DEFAULTS = yaml.safe_load(_CASES_FILE.read_text(encoding="utf-8")) or {}
if not CASE_DEFAULTS:
raise RuntimeError(f"Empty or missing case definitions in {_CASES_FILE}")
def _build_parser() -> FlexibleArgumentParser:
p = FlexibleArgumentParser(description="Offline Ming-omni-tts example")
p.add_argument("--model", default=_DEFAULT_MODEL, help="Model name or local path")
p.add_argument("--deploy-config", default=None, help="Deploy config YAML; auto-selected when omitted")
p.add_argument("--case", choices=sorted(CASE_DEFAULTS), default="style", help="Built-in demo case")
p.add_argument("--text", default=None, help="Override case text")
p.add_argument("--prompt", default=None, help="Override the system prompt prefix")
p.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
p.add_argument(
"--instruction-json", default=None, help='Structured Ming instruction JSON, e.g. \'{"方言":"广粤话"}\''
)
p.add_argument("--ref-audio", default=None, help="Single reference audio path for cloning-style cases")
p.add_argument("--ref-audio-paths", nargs="+", default=None, help="Multiple reference audio paths (podcast)")
p.add_argument("--ref-text", default=None, help="Reference transcript; required for zero_shot")
p.add_argument("--speaker-embedding", default=None, help="Path to a JSON speaker embedding file")
p.add_argument(
"--extract-speaker-embeddings",
action="store_true",
help="Extract speaker embeddings from ref audio via campplus.onnx",
)
p.add_argument("--max-decode-steps", type=int, default=None, help="Override ming_max_decode_steps")
p.add_argument("--output-dir", default="output_audio", help="Directory for output wav files")
p.add_argument("--output-name", default=None, help="Output wav filename")
p.add_argument("--num-prompts", type=int, default=1, help="Repeat the same prompt N times")
p.add_argument("--streaming", action="store_true", help="Use AsyncOmni with async_chunk streaming")
p.add_argument("--trust-remote-code", action="store_true")
p.add_argument("--enforce-eager", action="store_true")
p.add_argument(
"--log-stats", "--enable-stats", dest="log_stats", action="store_true", help="Enable Omni stats logging"
)
p.add_argument("--stats-log-file", default=None, help="Path for the Omni stats log file")
p.add_argument("--metadata-json", default=None, help="Path for the run manifest JSON")
p.add_argument("--stage-init-timeout", type=int, default=300, help="Per-stage init timeout (s)")
p.add_argument("--init-timeout", type=int, default=600, help="Total init timeout (s)")
p.add_argument("--batch-timeout", type=int, default=5, help="Batch timeout (s)")
p.add_argument("--shm-threshold-bytes", type=int, default=65536)
p.add_argument("--worker-backend", default="multi_process", choices=["multi_process", "ray"])
p.add_argument("--ray-address", default=None, help="Ray cluster address (--worker-backend ray)")
return p
def _finalize_args(args) -> None:
if args.instructions is not None and args.instruction_json is not None:
raise RuntimeError("Use either --instructions or --instruction-json, not both")
if args.num_prompts < 1:
raise RuntimeError("--num-prompts must be at least 1")
if args.streaming and args.num_prompts != 1:
raise RuntimeError("--streaming currently supports exactly one prompt")
if args.deploy_config is None:
args.deploy_config = _DEFAULT_DEPLOY_CONFIG
def _load_waveform(path: str) -> torch.Tensor:
samples, sr = sf.read(path, dtype="float32")
wav = torch.as_tensor(samples, dtype=torch.float32)
if wav.ndim == 2:
wav = wav.mean(dim=1)
wav = wav.reshape(1, -1)
if int(sr) != SAMPLE_RATE:
wav = torchaudio.functional.resample(wav, int(sr), SAMPLE_RATE)
return wav
def _load_speaker_embedding(path: str) -> torch.Tensor:
return torch.as_tensor(json.loads(Path(path).read_text(encoding="utf-8")), dtype=torch.float32)
def _ref_audio_paths(args) -> list[str]:
if args.ref_audio is not None and args.ref_audio_paths is not None:
raise RuntimeError("Use either --ref-audio or --ref-audio-paths, not both")
if args.ref_audio_paths is not None:
return list(args.ref_audio_paths)
return [args.ref_audio] if args.ref_audio else []
def _resolve_reference_inputs(args, case: dict, paths: list[str]):
required = int(case.get("requires_ref_audio_count", 0))
if required > 0 and len(paths) < required:
raise RuntimeError(f"Case '{args.case}' needs {required} ref audio paths via --ref-audio-paths")
if required <= 0 and case.get("requires_ref_audio") and not paths:
raise RuntimeError(f"--ref-audio required for case '{args.case}'")
if not paths:
return None
wavs = [_load_waveform(p) for p in paths]
return wavs[0] if len(wavs) == 1 else wavs
def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
if args.speaker_embedding:
return _load_speaker_embedding(args.speaker_embedding)
if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
return None
embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
if not embs:
raise RuntimeError("Speaker extraction produced no embeddings")
return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)
def _build_prompt_payload(tokenizer, args):
case = CASE_DEFAULTS[args.case]
paths = _ref_audio_paths(args)
prompt_text = args.ref_text if args.ref_text is not None else case.get("prompt_text")
if case.get("requires_ref_text") and not prompt_text:
raise RuntimeError(f"--ref-text required for case '{args.case}'")
prompt_waveform = _resolve_reference_inputs(args, case, paths) if prompt_text is not None else None
speaker_embedding = _resolve_speaker_embedding(args, case, paths)
use_zero_spk_emb = bool(case.get("use_zero_spk_emb")) and prompt_waveform is None and speaker_embedding is None
runtime_controls = {KEY_MAX_DECODE_STEPS: args.max_decode_steps or case["max_decode_steps"]}
for key, field in [(KEY_CFG, "cfg"), (KEY_SIGMA, "sigma"), (KEY_TEMPERATURE, "temperature")]:
if field in case:
runtime_controls[key] = case[field]
instruction = (
json.loads(args.instruction_json) if args.instruction_json else (args.instructions or case.get("instruction"))
)
return build_ming_dense_prompt(
tokenizer,
prompt=args.prompt or case["prompt"],
text=args.text or case["text"],
runtime_controls=runtime_controls,
instruction=instruction,
prompt_text=prompt_text,
prompt_waveform=prompt_waveform,
speaker_embedding=speaker_embedding,
use_zero_spk_emb=use_zero_spk_emb,
)
def main():
args = _build_parser().parse_args()
_finalize_args(args)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False)
prompt_payload = _build_prompt_payload(tokenizer, args)
case = CASE_DEFAULTS[args.case]
sampling_params_list = build_sampling_params(args.max_decode_steps or case["max_decode_steps"])
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
stats_log_file = resolve_stats_log_file(args)
summaries = run_generation(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
metadata_json = resolve_metadata_json(args)
manifest = build_manifest(args, prompt_payload, stats_log_file, summaries)
if metadata_json is not None:
Path(metadata_json).write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
print(f"Saved run manifest to {metadata_json}")
if __name__ == "__main__":
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
main()
@@ -1,277 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import json
import time
import uuid
import wave
from pathlib import Path
import torch
from vllm import SamplingParams
from vllm_omni import AsyncOmni, Omni
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
KEY_SPEAKER_EMBEDDING,
SAMPLE_RATE,
TEXT_EOS_TOKEN_ID,
)
def coerce_audio_tensor(audio, *, async_chunk: bool) -> torch.Tensor:
if isinstance(audio, list):
parts = []
for item in audio:
tensor = torch.as_tensor(item, dtype=torch.float32).reshape(-1)
if tensor.numel() > 0:
parts.append(tensor)
if not parts:
return torch.zeros((0,), dtype=torch.float32)
return torch.cat(parts, dim=0)
return torch.as_tensor(audio, dtype=torch.float32).reshape(-1)
def resolve_sr(sr) -> int:
if isinstance(sr, list):
sr = sr[-1]
if hasattr(sr, "item"):
return int(sr.item())
return int(sr)
def extract_sample_rate(multimodal_output: dict) -> int:
sr = multimodal_output.get("sr")
if sr is None:
raise RuntimeError("Expected multimodal_output['sr']")
return resolve_sr(sr)
def write_wav(path: str, audio: torch.Tensor, sample_rate: int) -> None:
audio = audio.clamp(-1.0, 1.0)
pcm16 = (audio * 32767.0).round().to(torch.int16).cpu().numpy()
with wave.open(path, "wb") as wav_file:
wav_file.setnchannels(1)
wav_file.setsampwidth(2)
wav_file.setframerate(int(sample_rate))
wav_file.writeframes(pcm16.tobytes())
def request_index(request_id: str | None, fallback: int) -> int:
try:
return int(request_id)
except (TypeError, ValueError):
if isinstance(request_id, str):
head = request_id.split("_", 1)[0]
if head.isdigit():
return int(head)
return fallback
def audio_summary(audio: torch.Tensor, sample_rate: int) -> dict:
waveform = audio.detach().cpu().reshape(-1).to(torch.float32)
return {
"sample_rate": int(sample_rate),
"num_samples": int(waveform.numel()),
"duration_seconds": float(waveform.numel()) / float(sample_rate),
"max_abs_amplitude": float(waveform.abs().max().item()) if waveform.numel() > 0 else 0.0,
}
def resolve_output_name(output_name: str | None, case: str, index: int, total: int) -> str:
if total == 1:
return output_name or f"ming_{case}.wav"
base = Path(output_name or f"ming_{case}.wav")
return f"{base.stem}_{index:05d}{base.suffix or '.wav'}"
def resolve_stats_log_file(args) -> str | None:
if not args.log_stats:
return None
if args.stats_log_file:
return args.stats_log_file
base = Path(args.output_name or f"ming_{args.case}.wav").stem
return str(Path(args.output_dir) / f"{base}_pipeline.log")
def resolve_metadata_json(args) -> str | None:
if args.metadata_json:
return args.metadata_json
if args.log_stats:
base = Path(args.output_name or f"ming_{args.case}.wav").stem
return str(Path(args.output_dir) / f"{base}_manifest.json")
return None
def build_manifest(args, prompt_payload, stats_log_file: str | None, outputs: list[dict]) -> dict:
additional_information = {}
if isinstance(prompt_payload, dict):
additional_information = dict(prompt_payload.get("additional_information", {}))
return {
"model": args.model,
"case": args.case,
"streaming": bool(args.streaming),
"deploy_config": args.deploy_config,
"enforce_eager": bool(args.enforce_eager),
"num_prompts": int(args.num_prompts),
"log_stats": bool(args.log_stats),
"stats_log_file": stats_log_file,
"prompt_text": additional_information.get("prompt_text"),
"instruction": additional_information.get("instruction"),
"speaker_embedding_shape": (
list(additional_information[KEY_SPEAKER_EMBEDDING].shape)
if KEY_SPEAKER_EMBEDDING in additional_information
and hasattr(additional_information[KEY_SPEAKER_EMBEDDING], "shape")
else None
),
"outputs": outputs,
"generated_at_unix": time.time(),
}
def build_engine_kwargs(args, stats_log_file: str | None) -> dict:
kwargs = {
"model": args.model,
"deploy_config": args.deploy_config,
"enforce_eager": args.enforce_eager,
"trust_remote_code": args.trust_remote_code,
"log_stats": args.log_stats,
"stage_init_timeout": args.stage_init_timeout,
"init_timeout": args.init_timeout,
"batch_timeout": args.batch_timeout,
"shm_threshold_bytes": args.shm_threshold_bytes,
"worker_backend": args.worker_backend,
}
if stats_log_file is not None:
kwargs["log_file"] = stats_log_file
if args.ray_address is not None:
kwargs["ray_address"] = args.ray_address
return kwargs
def build_sampling_params(max_decode_steps: int) -> list[SamplingParams]:
return [
SamplingParams(
temperature=0.0,
max_tokens=max_decode_steps + 1,
stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
),
SamplingParams(temperature=0.0, max_tokens=1),
]
async def run_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
engine = AsyncOmni(**build_engine_kwargs(args, stats_log_file))
try:
all_audio_chunks = []
accumulated_samples = 0
chunk_idx = 0
start_time = time.time()
chunk_times = []
ttfp_seconds = None
final_stage_output = None
async for stage_output in engine.generate(
prompt=prompt_payload,
request_id=str(uuid.uuid4()),
sampling_params_list=sampling_params_list,
):
final_stage_output = stage_output
multimodal_output = stage_output.multimodal_output or {}
audio = multimodal_output.get("audio")
if audio is None:
continue
finished = stage_output.finished
if isinstance(audio, torch.Tensor):
if finished:
audio_chunk = audio[accumulated_samples:].float().detach().cpu()
else:
audio_chunk = audio.float().detach().cpu()
elif isinstance(audio, list):
audio_chunk = torch.as_tensor(audio[chunk_idx], dtype=torch.float32).reshape(-1).cpu()
else:
audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
accumulated_samples += int(audio_chunk.numel())
chunk_idx += 1
if audio_chunk.numel() > 0:
now = time.time()
if ttfp_seconds is None:
ttfp_seconds = now - start_time
chunk_times.append(now)
all_audio_chunks.append(audio_chunk)
if not all_audio_chunks:
raise RuntimeError("Streaming Ming example produced no audio chunks")
waveform = torch.cat(all_audio_chunks, dim=0)
output_name = resolve_output_name(args.output_name, args.case, 0, 1)
output_path = str(output_dir / output_name)
write_wav(output_path, waveform, SAMPLE_RATE)
summary = {
"request_id": getattr(final_stage_output, "request_id", None),
"stage_id": getattr(final_stage_output, "stage_id", None),
"output_path": output_path,
"stage_durations": getattr(final_stage_output, "stage_durations", {}),
"peak_memory_mb": getattr(final_stage_output, "peak_memory_mb", 0.0),
"ttfp_seconds": ttfp_seconds,
"mean_inter_chunk_seconds": (
sum(t1 - t0 for t0, t1 in zip(chunk_times, chunk_times[1:])) / (len(chunk_times) - 1)
if len(chunk_times) > 1
else None
),
}
summary.update(audio_summary(waveform, SAMPLE_RATE))
print(f"Saved streaming output to {output_path}")
print(json.dumps(summary, ensure_ascii=False, indent=2))
return [summary]
finally:
engine.shutdown()
def run_non_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
engine = Omni(**build_engine_kwargs(args, stats_log_file))
try:
outputs = engine.generate(
prompts=[prompt_payload for _ in range(args.num_prompts)],
sampling_params_list=sampling_params_list,
py_generator=False,
)
summaries = []
for fallback_index, output in enumerate(outputs):
if output.final_output_type != "audio":
continue
multimodal_output = output.multimodal_output or {}
waveform = coerce_audio_tensor(multimodal_output.get("audio"), async_chunk=False)
sample_rate = extract_sample_rate(multimodal_output)
output_name = resolve_output_name(
args.output_name,
args.case,
request_index(output.request_id, fallback_index),
args.num_prompts,
)
output_path = str(output_dir / output_name)
write_wav(output_path, waveform, sample_rate)
summary = {
"request_id": output.request_id,
"stage_id": output.stage_id,
"output_path": output_path,
"stage_durations": output.stage_durations,
"peak_memory_mb": output.peak_memory_mb,
}
summary.update(audio_summary(waveform, sample_rate))
summaries.append(summary)
print(f"Saved output to {output_path}")
print(json.dumps(summary, ensure_ascii=False, indent=2))
if not summaries:
raise RuntimeError("Non-streaming Ming example produced no audio outputs")
return summaries
finally:
engine.close()
def run_generation(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
if args.streaming:
return asyncio.run(run_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file))
return run_non_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
@@ -1,377 +0,0 @@
"""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)
Usage
-----
# Standard TTS with voice cloning:
python end2end.py \\
--model OpenMOSS-Team/MOSS-TTS \\
--text "Hello, this is a MOSS-TTS test." \\
--ref-audio path/to/reference.wav \\
--output output.wav
"""
from __future__ import annotations
import argparse
import gc
import sys
import soundfile as sf
import torch
def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
"""Load and resample reference audio to (1, T) at target_sr."""
data, sr = sf.read(path, always_2d=True) # (T, C)
wav = torch.from_numpy(data.T.astype("float32")) # (C, T)
if wav.shape[0] > 1:
wav = wav.mean(dim=0, keepdim=True)
if sr != target_sr:
import torchaudio
wav = torchaudio.functional.resample(wav, sr, target_sr)
return wav # (1, T)
def _build_unified_codes(
model_id: str,
text: str,
ref_audio: torch.Tensor | None,
*,
instruction: str | None = None,
tokens: int | None = None,
quality: str | None = None,
sound_event: str | None = None,
ambient_sound: str | None = None,
language: str | None = None,
) -> tuple[list[int], torch.Tensor, int]:
"""Build the upstream MossTTSDelay unified-codes prompt.
Returns (text_token_ids, audio_codes_LxNQ, n_vq) where
``text_token_ids`` is the prefill prompt for the talker and
``audio_codes_LxNQ`` is the matching delay-pattern audio code grid
(``audio_pad_code`` everywhere except inside the reference-audio block).
"""
from transformers import AutoProcessor
proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
reference: list[torch.Tensor] | None = None
if ref_audio is not None:
# Pre-encode so the processor doesn't need torchaudio.load (which
# now requires torchcodec on recent torchaudio releases).
# encode_audios_from_wav expects (C, T) per item.
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
codes_list = proc.encode_audios_from_wav([wav_2d], sampling_rate=24000, n_vq=proc.model_config.n_vq)
reference = [codes_list[0]]
user_msg_kwargs: dict = {"text": text}
if reference is not None:
user_msg_kwargs["reference"] = reference
if instruction is not None:
user_msg_kwargs["instruction"] = instruction
if tokens is not None:
user_msg_kwargs["tokens"] = tokens
if quality is not None:
user_msg_kwargs["quality"] = quality
if sound_event is not None:
user_msg_kwargs["sound_event"] = sound_event
if ambient_sound is not None:
user_msg_kwargs["ambient_sound"] = ambient_sound
if language is not None:
user_msg_kwargs["language"] = language
user_msg = proc.build_user_message(**user_msg_kwargs)
batch = proc(conversations=[[user_msg]], mode="generation")
unified = batch["input_ids"][0] # (L, 1+n_vq)
n_vq = int(unified.shape[-1] - 1)
text_ids = unified[:, 0].tolist()
audio_codes = unified[:, 1:].contiguous().to(torch.int64) # (L, n_vq)
# Free the processor (incl. ~1.6B-param audio_tokenizer on CPU) before
# we instantiate Omni — we only needed it for the one-shot encode.
del proc
gc.collect()
return text_ids, audio_codes, n_vq
def _build_realtime_prompt(
model_id: str,
text: str,
ref_audio: torch.Tensor | None,
) -> tuple[list[int], torch.Tensor, int]:
"""Build the prefill prompt for MOSS-TTS-Realtime.
Realtime uses a different prompt format than the delay variants: a system
prompt + (optional) voice-clone context with the encoded reference audio
inlined, followed by the user text and an assistant header. The format is
a ``(L, channels+1=17)`` int grid where col 0 is text/special tokens and
cols 1..16 are RVQ codebook entries (``audio_channel_pad`` outside the
ref-audio block).
The upstream ``MossTTSRealtimeProcessor`` is shipped via
``trust_remote_code`` but is not auto-discovered by ``AutoProcessor``
(no ``processor_config.json``) — we import its module directly from the
snapshot directory and use its ``make_ensemble``/``make_user_prompt``.
"""
import importlib.util
import sys as _sys
from pathlib import Path
import numpy as np
import torch
from huggingface_hub import snapshot_download
from transformers import AutoModel, AutoTokenizer
snap_dir = Path(snapshot_download(repo_id=model_id))
proc_module_path = snap_dir / "processing_mossttsrealtime.py"
if not proc_module_path.exists():
raise FileNotFoundError(f"realtime processor module missing at {proc_module_path}")
# Load the processor module without polluting sys.path globally.
spec = importlib.util.spec_from_file_location("_moss_tts_realtime_proc", proc_module_path)
mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type]
_sys.modules[spec.name] = mod # type: ignore[union-attr]
spec.loader.exec_module(mod) # type: ignore[union-attr]
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
# Encode reference audio via the standalone MOSS Audio Tokenizer (CPU is
# fine — it runs once per request).
audio_tokens = None
if ref_audio is not None:
codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).eval()
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
with torch.no_grad():
enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
# enc.audio_codes: (NQ, B, T); we want (T, NQ) numpy.
codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy() # (T, 16)
del codec
gc.collect()
# Build the inference-style prefill (mirrors upstream
# ``MossTTSRealtimeInference._build_prefill_batch``):
# system_prompt [+ voice_clone_block]
# <|im_start|>assistant\n
# <text_tokens...> (col 0 = text, cols 1..16 = audio_channel_pad)
# <last text token> (col 1 = audio_bos at the last text position)
# Generation continues from here; the local transformer emits real audio
# codes for codebook 0 directly (no BOS frame at gen_step=0).
system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens) # (L_sys, 17)
# Append "<|im_start|>assistant\n" header (the snapshot's make_ensemble
# omits it; inferencer.py's version includes it. We add it here.)
assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
assistant_grid[:, 0] = assistant_header_ids
# Append the FIRST 12 text tokens with audio_bos (1025) at the last
# included text position (audio column 0 == grid col 1). The remaining
# text tokens are fed one-per-step during generation via the talker's
# compute_logits → vLLM sampler pipeline (see ``ids.all`` in the
# additional_information). This matches upstream
# ``MossTTSRealtimeInference._build_prefill_batch`` (prefill_max_text=12).
text_ids_only = tokenizer.encode(text, add_special_tokens=False)
if not text_ids_only:
raise ValueError("Empty text after tokenisation")
PREFILL_MAX_TEXT = 12
cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
text_grid[:, 0] = text_ids_only[:cur_len]
text_grid[-1, 1] = 1025 # audio_bos_token at last *prefilled* text token
grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
text_ids = grid[:, 0].tolist()
audio_codes = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16)
# ``remaining_text_ids`` will be fed one-per-step by the talker.
remaining_text_ids = list(text_ids_only[cur_len:])
return text_ids, audio_codes, 16, remaining_text_ids
def run_tts(args: argparse.Namespace) -> None:
from vllm import SamplingParams
from vllm_omni import Omni
deploy_map = {
"OpenMOSS-Team/MOSS-TTS": "moss_tts",
"OpenMOSS-Team/MOSS-TTSD-v1.0": "moss_ttsd",
"OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
"OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
"OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
}
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
ref_audio = None
builder_kwargs: dict = {}
remaining_text_ids: list[int] | None = None
if is_realtime:
# MOSS-TTS-Realtime uses a different prompt format and has its own
# per-step depth transformer; the delay-style ``MossTTSDelayProcessor``
# path doesn't apply.
if not args.ref_audio:
sys.exit("MOSS-TTS-Realtime currently requires --ref-audio (voice clone).")
ref_audio = _load_ref_audio(args.ref_audio)
text = args.text or ""
text_ids, audio_codes, n_vq, remaining_text_ids = _build_realtime_prompt(args.model, text, ref_audio)
print(
f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block (realtime); "
f"{len(remaining_text_ids)} text tokens to stream"
)
elif is_sound_effect:
if not args.ambient_sound:
sys.exit("MOSS-SoundEffect requires --ambient-sound")
builder_kwargs["ambient_sound"] = args.ambient_sound
# MOSS-SoundEffect emits ~12.5 codec frames per second; ``tokens``
# tells the AR loop the target duration.
if args.duration_tokens:
builder_kwargs["tokens"] = args.duration_tokens
elif args.duration_seconds:
builder_kwargs["tokens"] = max(1, int(float(args.duration_seconds) * 12.5))
text = args.text or ""
elif is_voice_gen:
# MOSS-VoiceGenerator needs an instruction (voice description) plus
# the text to synthesise; no reference audio.
if not args.instruction:
sys.exit("MOSS-VoiceGenerator requires --instruction (voice description)")
builder_kwargs["instruction"] = args.instruction
text = args.text or ""
else:
if not args.ref_audio:
sys.exit("Voice cloning requires --ref-audio")
ref_audio = _load_ref_audio(args.ref_audio)
text = args.text or ""
if not is_realtime:
text_ids, audio_codes, n_vq = _build_unified_codes(args.model, text, ref_audio, **builder_kwargs)
print(f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block")
omni = Omni(model=args.model, deploy_config=deploy_config, stage_init_timeout=600)
talker_sp = SamplingParams(
temperature=1.7,
top_p=0.8,
top_k=25,
max_tokens=2048,
seed=42,
detokenize=False,
)
codec_sp = SamplingParams(
temperature=0.0,
top_p=1.0,
top_k=-1,
max_tokens=65536,
seed=42,
detokenize=True,
)
additional_info: dict = {
# Reference-audio codes consumed by the talker prefill (additive
# embedding); shape is (L, n_vq) where pad_code marks non-audio
# positions. The OmniPayload schema reserves ``codes.ref`` for
# exactly this kind of conditioning input.
"codes": {"ref": audio_codes},
}
if remaining_text_ids:
# Realtime variant: feed the remaining text tokens (everything past
# the first ``PREFILL_MAX_TEXT``) one-per-step during decode via the
# talker's compute_logits → vLLM sampler one-hot trick.
additional_info["ids"] = {"all": remaining_text_ids}
request = {
"prompt_token_ids": text_ids,
"additional_information": additional_info,
}
import time as _time
chunks: list[torch.Tensor] = []
t_gen_start = _time.monotonic()
t_first_audio: float | None = None
for stage_outputs in omni.generate(request, [talker_sp, codec_sp]):
# Each yielded ``stage_outputs`` is an OmniRequestOutput; the codec
# stage attaches its waveform under ``multimodal_output["model_outputs"]``.
# Avoid ``a or b`` when ``a`` may be a tensor (truthiness raises).
mm = stage_outputs.multimodal_output or {}
wav = mm.get("audio")
if wav is None:
wav = mm.get("model_outputs")
if wav is None:
continue
if isinstance(wav, list):
non_empty = [w.reshape(-1) 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))
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}")
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}")
def main() -> None:
parser = argparse.ArgumentParser(description="MOSS-TTS offline inference example")
parser.add_argument(
"--model",
default="OpenMOSS-Team/MOSS-TTS",
help="HuggingFace model ID",
)
parser.add_argument("--text", default="", help="Text to synthesise (TTS variants)")
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning")
parser.add_argument(
"--ambient-sound",
default=None,
help="Sound description for MOSS-SoundEffect",
)
parser.add_argument(
"--duration-tokens",
type=int,
default=None,
help="Target duration in tokens (1 s ≈ 12.5 tokens) for MOSS-SoundEffect",
)
parser.add_argument(
"--duration-seconds",
type=float,
default=None,
help="Target duration in seconds (converted via ~12.5 tokens/s) for MOSS-SoundEffect",
)
parser.add_argument(
"--instruction",
default=None,
help="Voice description for MOSS-VoiceGenerator (e.g. 'a young woman with a clear voice')",
)
parser.add_argument("--output", default="output.wav", help="Output WAV path")
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
@@ -38,8 +38,7 @@ from pathlib import Path
import soundfile as sf
import torch
from vllm import SamplingParams
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
from vllm.utils.argparse_utils import FlexibleArgumentParser
# Prevent multiprocessing from re-importing CUDA in the wrong context.
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
@@ -156,7 +155,7 @@ def main(args) -> None:
def parse_args():
parser = TrackingArgumentParser(description="MOSS-TTS-Nano offline inference")
parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
parser.add_argument(
"--ref-audio",
@@ -21,6 +21,7 @@ import os
import numpy as np
import soundfile as sf
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -85,6 +86,7 @@ def run_e2e():
default=None,
help="Random seed for generation",
)
nullify_stage_engine_defaults(parser)
args = parser.parse_args()
if not os.path.exists(args.stage_config):
@@ -15,8 +15,10 @@ import torch
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
logger = logging.getLogger(__name__)
@@ -433,7 +435,7 @@ async def main_streaming(args):
def parse_args():
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--query-type",
"-q",
@@ -546,6 +548,7 @@ def parse_args():
help="Number of prompts per batch (default: 1, sequential).",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -12,16 +12,17 @@ from pathlib import Path
import soundfile as sf
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
REPO_ROOT = Path(__file__).resolve().parents[4]
SAMPLE_RATE = 48_000
def parse_args():
parser = TrackingArgumentParser(description="Offline VoxCPM2 native AR inference")
parser = FlexibleArgumentParser(description="Offline VoxCPM2 native AR inference")
parser.add_argument(
"--model",
type=str,
@@ -59,6 +60,7 @@ def parse_args():
default=None,
help="Optional transcript of --ref-audio (enables continuation mode).",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -27,10 +27,11 @@ except ImportError:
"pip install -e ."
)
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
@@ -212,7 +213,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
output_audio_dur = 0.0
for batch_idx, o in enumerate(outputs):
audio_tensor = torch.cat((o.multimodal_output["audio"],))
audio_tensor = torch.cat(o.multimodal_output["audio"])
audio_array = audio_tensor.tolist()
output_audio_dur += float(len(audio_array)) / 24000
if args.write_audio:
@@ -229,7 +230,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
def parse_args() -> Namespace:
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
parser.add_argument(
"--model",
type=str,
@@ -312,6 +313,7 @@ def parse_args() -> Namespace:
default=None,
help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
-63
View File
@@ -68,67 +68,6 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size
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.
@@ -384,8 +323,6 @@ python openai_chat_client.py \
| `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`) |
-132
View File
@@ -1,132 +0,0 @@
# DreamZero
DreamZero is the robot-policy serving example for the OpenPI-compatible `/v1/realtime/robot/openpi` endpoint.
## Supported checkpoints
| Model | HuggingFace repo |
|---|---|
| DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` |
## Quick start
### Start the server
```bash
bash examples/online_serving/dreamzero/run_server.sh
```
By default this launches:
```bash
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
--served-model-name dreamzero-droid \
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
--enforce-eager --disable-log-stats
```
Override `MODEL`, `PORT`, `HOST`, `DEPLOY_CONFIG`, or `SERVED_MODEL_NAME` through the script environment if needed.
### Install example dependencies
The core `pip install -e .` setup does not include the extra packages used by the DreamZero example scripts.
- `openpi_client.py`:
`openpi-client`, `websockets`, `opencv-python`
- `droid_sim_eval_client.py`:
`mediapy`, `websockets`, `openpi-client`
The DROID sim-eval script also needs an Isaac Lab environment that provides `isaaclab`, `isaaclab_tasks`, `sim_evals`, and `gymnasium`.
If you run the DROID client on Python < 3.12, also install `typing-extensions`.
### Configure TP and CFG parallelism
The bundled DreamZero configs intentionally keep only:
| Config | Purpose |
|---|---|
| `vllm_omni/deploy/dreamzero.yaml` | Default TP=1, CFG parallel disabled |
| `vllm_omni/deploy/dreamzero_tp1_cfg2.yaml` | TP=1, CFG parallel size=2 |
For other topologies, use CLI parallelism flags and update stage 0 `devices` with `--stage-overrides`. The number of listed devices must match `tensor_parallel_size * cfg_parallel_size`.
TP=2 with CFG parallel disabled:
```bash
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
--served-model-name dreamzero-droid \
--deploy-config vllm_omni/deploy/dreamzero.yaml \
--tensor-parallel-size 2 \
--stage-overrides '{"0": {"devices": "0,1"}}' \
--enforce-eager --disable-log-stats
```
TP=2 with CFG parallel size=2:
```bash
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
--served-model-name dreamzero-droid \
--deploy-config vllm_omni/deploy/dreamzero.yaml \
--tensor-parallel-size 2 \
--cfg-parallel-size 2 \
--stage-overrides '{"0": {"devices": "0,1,2,3"}}' \
--enforce-eager --disable-log-stats
```
### Download example assets
The OpenPI client and DROID sim-eval example expect the three camera MP4 files in `outputs/dreamzero/assets`.
```bash
hf download YangshenDeng/vllm-omni-dreamzero-assets --repo-type dataset --local-dir outputs/dreamzero/assets
```
### Run the OpenPI client
```bash
python examples/online_serving/dreamzero/openpi_client.py \
--host 127.0.0.1 \
--port 8000 \
--video-dir outputs/dreamzero/assets
```
This client uses downloaded example videos and talks to the OpenPI websocket server.
### Run DROID sim eval
```bash
${ISAACLAB_LAUNCHER} -p examples/online_serving/dreamzero/droid_sim_eval_client.py \
--host 127.0.0.1 \
--port 8000 \
--scene 1 \
--episodes 1 \
--headless \
--device cuda:0
```
Set `ISAACLAB_LAUNCHER=path/to/isaaclab.sh` from the vLLM-Omni repository root before running the command.
This launches Isaac Lab / sim-evals and runs the DROID benchmark loop against the same websocket endpoint.
### Export comparison videos offline
```bash
python examples/offline_inference/dreamzero/export_prediction_video.py \
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
--save-input-video \
--save-gif
```
The export script uses local `Omni` inference, not the websocket server. It writes the input rollout video and the predicted output video artifacts for side-by-side inspection.
### MolmoSpace demo
```bash
python examples/online_serving/dreamzero/molmospace_dreamzero_eval_demo.py \
--host 127.0.0.1 \
--port 8000 \
--benchmark_dir /path/to/benchmark \
--output_dir /path/to/output
```
This demo adapts DreamZero to the MolmoSpace-style remote policy eval loop.
@@ -1,819 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Run one or more DROID sim-eval rollouts against the vLLM DreamZero server.
This script is the vLLM/OpenPI adaptation of the upstream DreamZero
`eval_utils/run_sim_eval.py` sim-eval client.
Behavior intentionally kept close to upstream:
- same DROID observation extraction (`external_cam`, `external_cam_2`,
`wrist_cam`, joint position, gripper position)
- same resize-with-pad preprocessing to `(180, 320)`
- same `open_loop_horizon=8`
- same gripper binarization rule (`> 0.5 -> 1`, else `0`)
- same per-scene language prompts
Unlike upstream DreamZero, vLLM serves the compatible websocket policy endpoint
at `/v1/realtime/robot/openpi`, so this script includes the path suffix in the
client URI.
Run this script through Isaac Lab's launcher from the vLLM-Omni repository
root, for example:
"${ISAACLAB_LAUNCHER}" -p \
examples/online_serving/dreamzero/droid_sim_eval_client.py \
--host 127.0.0.1 \
--port 8000 \
--scene 1 \
--episodes 1 \
--headless \
--device cuda:1
"""
from __future__ import annotations
import argparse
import json
import logging
import sys
import time
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
import numpy as np
import torch
try:
import mediapy
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `mediapy`.") from exc
try:
from typing import override
except ImportError:
try:
from typing_extensions import override
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `typing-extensions` on Python < 3.12.") from exc
try:
import websockets.sync.client
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `websockets`.") from exc
# NOTE:
# This directory already contains a local file named `openpi_client.py`.
# However, what we want here is the *installed* `openpi_client` package from
# upstream OpenPI, not the sibling example file. When a script is executed
# directly, Python often puts the script directory at `sys.path[0]`, which
# would cause `import openpi_client` to resolve to the local example file and
# create a circular import.
#
# To avoid that ambiguity, temporarily remove the current example directory
# from the front of `sys.path`, import the real package, and then restore the
# path afterwards.
example_dir = str(Path(__file__).resolve().parent)
removed_path = False
if sys.path and sys.path[0] == example_dir:
sys.path.pop(0)
removed_path = True
try:
from openpi_client import image_tools, msgpack_numpy
from openpi_client.base_policy import BasePolicy
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires the `openpi-client` package.") from exc
finally:
if removed_path:
sys.path.insert(0, example_dir)
# -----------------------------------------------------------------------------
# Constants
# -----------------------------------------------------------------------------
#
# These values intentionally mirror the upstream DreamZero sim-eval client
# where possible. The important distinction is:
#
# - ACTION_HORIZON = 24
# The model returns 24 future actions per inference call.
# - DEFAULT_OPEN_LOOP_HORIZON = 8
# The sim client only executes the first 8 actions locally before asking
# the server to replan from a fresh observation.
#
# So a single server call predicts 24x8 actions, but the rollout consumes only
# 8 of them before replanning.
PING_INTERVAL_SECS = 60
PING_TIMEOUT_SECS = 600
DEFAULT_PATH = "/v1/realtime/robot/openpi"
DEFAULT_OPEN_LOOP_HORIZON = 8
ACTION_HORIZON = 24
ACTION_DIM = 8
DEFAULT_OUTPUT_ROOT = Path("runs") / "dreamzero_sim_eval"
SCENE_PROMPTS = {
1: "put the cube in the bowl",
2: "pick up the can and put it in the mug",
3: "put the banana in the bin",
}
def _decode_action_response(response: bytes | str) -> np.ndarray:
if isinstance(response, str):
raise RuntimeError(f"Error in inference server:\n{response}")
decoded = msgpack_numpy.unpackb(response)
if isinstance(decoded, dict) and decoded.get("type") == "error":
message = decoded.get("message", decoded)
raise RuntimeError(f"Error in inference server:\n{message}")
return np.asarray(decoded, dtype=np.float32)
@dataclass(frozen=True)
class StepRecord:
"""One fully materialized rollout step for later JSON export.
The `episode_00.json` artifact is intended to be human-readable and
post-process-friendly. Instead of keeping raw tensors around, each step is
flattened into plain Python types so it can be serialized directly.
"""
# Index of this control step within the episode.
step_index: int
# Whether this step triggered a fresh model call (as opposed to reusing the
# cached open-loop chunk from the previous server response).
used_server_call: bool
# End-to-end latency of the server call that produced the current chunk.
# This is `None` on steps that only reuse cached actions.
chunk_latency_s: float | None
# The concrete 8-D action sent into the simulator at this step.
action: list[float]
# Observed 7-DoF arm joint positions before the next environment step.
joint_position: list[float]
# Observed gripper scalar before the next environment step.
gripper_position: list[float]
# Reward and termination signals directly returned by the simulator.
reward: float
terminated: bool
truncated: bool
# Optional scene object positions for downstream debugging / success
# heuristics. This may be empty if the environment does not expose them.
scene_objects: dict[str, list[float]]
class OpenPIWebsocketClientPolicy(BasePolicy):
"""Minimal websocket client for the DreamZero/OpenPI policy protocol.
Protocol shape:
- connect -> server immediately sends a metadata dict
- infer -> send msgpack observation, receive action chunk
- reset -> send msgpack reset command, receive confirmation string
This class intentionally stays very small because the more interesting
DreamZero-specific behavior lives one layer above, in
`DreamZeroJointPosClient`.
"""
def __init__(
self,
host: str = "127.0.0.1",
port: int = 8000,
path: str = DEFAULT_PATH,
) -> None:
# vLLM serves the robot endpoint under `/v1/realtime/robot/openpi`.
self._uri = f"ws://{host}:{port}{path}"
# Upstream protocol uses msgpack with numpy support, not JSON.
self._packer = msgpack_numpy.Packer()
# Connect immediately and cache the server handshake metadata.
self._ws, self._server_metadata = self._wait_for_server()
def get_server_metadata(self) -> dict[str, Any]:
"""Return a copy of the server handshake metadata."""
return dict(self._server_metadata)
def _wait_for_server(self):
"""Connect to the websocket server and read the initial metadata frame."""
logging.info("Connecting to %s", self._uri)
conn = websockets.sync.client.connect(
self._uri,
compression=None,
max_size=None,
ping_interval=PING_INTERVAL_SECS,
ping_timeout=PING_TIMEOUT_SECS,
)
metadata = msgpack_numpy.unpackb(conn.recv())
if not isinstance(metadata, dict):
raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
return conn, metadata
@override
def infer(self, obs: dict[str, Any]) -> np.ndarray:
"""Send an inference request and return the decoded action chunk."""
# Keep the upstream DreamZero/OpenPI convention that the request itself
# tells the server which logical endpoint is being called.
payload = dict(obs)
payload["endpoint"] = "infer"
self._ws.send(self._packer.pack(payload))
response = self._ws.recv()
return _decode_action_response(response)
@override
def reset(self, reset_info: dict[str, Any] | None = None) -> str:
"""Tell the server to reset its session-side state."""
payload = dict(reset_info or {})
payload["endpoint"] = "reset"
self._ws.send(self._packer.pack(payload))
response = self._ws.recv()
if isinstance(response, str):
return response
decoded = msgpack_numpy.unpackb(response)
if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
raise RuntimeError(f"Unexpected reset response: {decoded!r}")
return str(decoded["status"])
def close(self) -> None:
"""Close the websocket connection explicitly."""
self._ws.close()
class DreamZeroJointPosClient:
"""DROID sim-eval client that talks to the vLLM OpenPI websocket server.
This is the main compatibility layer between:
- Isaac Lab DROID observations (`obs["policy"][...]`)
- DreamZero/OpenPI websocket payloads
- local open-loop action reuse across several simulator steps
In other words:
simulator obs -> websocket request -> action chunk -> one action per step
"""
def __init__(
self,
remote_host: str = "127.0.0.1",
remote_port: int = 8000,
path: str = DEFAULT_PATH,
open_loop_horizon: int = DEFAULT_OPEN_LOOP_HORIZON,
) -> None:
# Low-level transport client.
self.client = OpenPIWebsocketClientPolicy(remote_host, remote_port, path=path)
# Number of actions to execute locally before replanning.
self.open_loop_horizon = open_loop_horizon
# Cursor into the currently cached action chunk.
self.actions_from_chunk_completed = 0
# Most recent `(ACTION_HORIZON, ACTION_DIM)` server response.
self.pred_action_chunk: np.ndarray | None = None
# Session id is part of the DreamZero serving contract. Changing it
# causes the server side to treat the rollout as a fresh episode.
self.session_id = str(uuid.uuid4())
# Simple runtime stats for reporting.
self.server_calls = 0
self.last_chunk_latency_s: float | None = None
self.last_used_server_call = False
def metadata(self) -> dict[str, Any]:
"""Expose the server metadata to callers / logs."""
return self.client.get_server_metadata()
def reset(self) -> str:
"""Reset local chunk state and remote session state.
Local reset:
- drop cached action chunk
- rewind chunk cursor
- allocate a fresh session id
Remote reset:
- send a websocket `reset` message so the server can clear any
request/session-side state it associates with this client
"""
self.actions_from_chunk_completed = 0
self.pred_action_chunk = None
self.session_id = str(uuid.uuid4())
self.last_chunk_latency_s = None
self.last_used_server_call = False
return self.client.reset({})
def infer(self, obs: dict[str, Any], instruction: str) -> dict[str, Any]:
"""Turn one simulator observation into one executable 8-D action.
Key behavior:
- call the server only when the local chunk cache is empty/exhausted
- otherwise, keep consuming the cached chunk open-loop
- always return exactly one 8-D action for the current simulator step
"""
# Convert Isaac Lab observation structure into a plain numpy-friendly
# record that is easier to serialize and visualize.
curr_obs = self._extract_observation(obs)
self.last_used_server_call = False
# Replan if:
# 1. this is the first step of a rollout / chunk
# 2. we already consumed `open_loop_horizon` actions from the current chunk
# 3. no cached chunk is currently available
if (
self.actions_from_chunk_completed == 0
or self.actions_from_chunk_completed >= self.open_loop_horizon
or self.pred_action_chunk is None
):
self.actions_from_chunk_completed = 0
# Build the exact DreamZero/OpenPI payload expected by the server.
#
# Notes:
# - images are resized/padded to the serving contract's 180x320
# - proprio is cast to float64 to match upstream client behavior
# - cartesian_position is currently unused by DreamZero DROID, so
# a dummy zero vector is sent for protocol completeness
request_data = {
"observation/exterior_image_0_left": image_tools.resize_with_pad(curr_obs["right_image"], 180, 320),
"observation/exterior_image_1_left": image_tools.resize_with_pad(curr_obs["left_image"], 180, 320),
"observation/wrist_image_left": image_tools.resize_with_pad(curr_obs["wrist_image"], 180, 320),
"observation/joint_position": curr_obs["joint_position"].astype(np.float64),
"observation/cartesian_position": np.zeros((6,), dtype=np.float64),
"observation/gripper_position": curr_obs["gripper_position"].astype(np.float64),
"prompt": instruction,
"session_id": self.session_id,
}
# Measure end-to-end server latency for this chunk request.
start = time.perf_counter()
actions = self.client.infer(request_data)
self.last_chunk_latency_s = time.perf_counter() - start
self.last_used_server_call = True
self.server_calls += 1
# DreamZero DROID serving is expected to return an action chunk with
# 24 future actions, each action being 8-D.
if actions.ndim != 2:
raise AssertionError(f"Expected 2D action array, got shape {actions.shape}")
if actions.shape != (ACTION_HORIZON, ACTION_DIM):
raise AssertionError(f"Expected action shape {(ACTION_HORIZON, ACTION_DIM)}, got {actions.shape}")
self.pred_action_chunk = actions
# Consume exactly one action row from the cached chunk for this
# simulator step.
action = np.array(self.pred_action_chunk[self.actions_from_chunk_completed], copy=True)
self.actions_from_chunk_completed += 1
# Upstream DreamZero sim-eval binarizes the gripper command.
action[-1] = 1.0 if action[-1].item() > 0.5 else 0.0
# Produce a human-friendly visualization strip for videos:
# right external | wrist | left external
img1 = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
img2 = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
img3 = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
viz = np.concatenate([img1, img2, img3], axis=1)
# Return both the executable action and auxiliary debug info.
return {
"action": action,
"viz": viz,
"joint_position": curr_obs["joint_position"],
"gripper_position": curr_obs["gripper_position"],
"used_server_call": self.last_used_server_call,
"chunk_latency_s": self.last_chunk_latency_s if self.last_used_server_call else None,
}
@staticmethod
def _extract_observation(obs_dict: dict[str, Any]) -> dict[str, np.ndarray]:
"""Extract the pieces DreamZero cares about from Isaac Lab observations.
`sim-evals` exposes camera frames and robot state inside the
`obs["policy"]` group. This helper converts those tensors into numpy
arrays so they can be fed into image preprocessing / websocket packing.
"""
policy = obs_dict["policy"]
# Isaac Lab stores camera observations as batched tensors; use env 0.
right_image = policy["external_cam"][0].clone().detach().cpu().numpy()
left_image = policy["external_cam_2"][0].clone().detach().cpu().numpy()
wrist_image = policy["wrist_cam"][0].clone().detach().cpu().numpy()
# Robot proprioception.
joint_position = policy["arm_joint_pos"].clone().detach().cpu().numpy()
gripper_position = policy["gripper_pos"].clone().detach().cpu().numpy()
return {
"right_image": right_image,
"left_image": left_image,
"wrist_image": wrist_image,
"joint_position": joint_position,
"gripper_position": gripper_position,
}
def _scene_instruction(scene: int) -> str:
"""Map a numeric scene id onto the fixed language prompt used for rollout."""
try:
return SCENE_PROMPTS[scene]
except KeyError as exc:
raise ValueError(f"Unsupported scene {scene}. Available scenes: {sorted(SCENE_PROMPTS)}") from exc
def _capture_scene_objects(env: Any) -> dict[str, list[float]]:
"""Best-effort extraction of scene object root positions.
This is only for debugging / reporting. The rollout logic itself does not
depend on these positions.
"""
objects: dict[str, list[float]] = {}
scene = getattr(env, "scene", None)
if scene is None:
return objects
# Skip non-task entities such as cameras, the robot, and lighting.
for name in scene.keys():
if name in {"robot", "external_cam", "external_cam_2", "wrist_cam", "sphere_light", "scene"}:
continue
entity = scene[name]
data = getattr(entity, "data", None)
root_pos_w = getattr(data, "root_pos_w", None)
if root_pos_w is None:
continue
# Convert tensors to plain lists for JSON serialization.
value = root_pos_w[0].detach().cpu().to(torch.float32).tolist()
objects[str(name)] = [float(x) for x in value]
return objects
def _maybe_infer_success(scene: int, final_objects: dict[str, list[float]]) -> dict[str, Any]:
"""Best-effort geometric heuristic.
The simulator itself does not expose a built-in success term; this function
provides a transparent fallback for human-readable reporting only.
"""
task_pairs = {
1: ("cube", "bowl"),
2: ("can", "mug"),
3: ("banana", "bin"),
}
source_name, target_name = task_pairs.get(scene, (None, None))
if source_name not in final_objects or target_name not in final_objects:
return {
"has_builtin_success": False,
"heuristic_success": None,
"reason": "scene object names unavailable for heuristic",
}
source = np.asarray(final_objects[source_name], dtype=np.float32)
target = np.asarray(final_objects[target_name], dtype=np.float32)
xy_distance = float(np.linalg.norm(source[:2] - target[:2]))
z_delta = float(source[2] - target[2])
heuristic_success = bool(xy_distance < 0.12 and z_delta > -0.08)
return {
"has_builtin_success": False,
"heuristic_success": heuristic_success,
"xy_distance": xy_distance,
"z_delta": z_delta,
"source_object": source_name,
"target_object": target_name,
"reason": "no built-in env success flag; using final object pose heuristic",
}
def _dump_json(path: Path, payload: dict[str, Any]) -> None:
"""Write a JSON file with stable UTF-8 formatting."""
path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
def _scalar_from_env_value(value: Any) -> float:
"""Normalize simulator scalar outputs into a plain Python float.
Isaac Lab / Gym values may come back as tensors, numpy arrays, tuples, or
direct Python scalars depending on the wrapper stack. Centralizing the
conversion here makes the rollout loop cleaner and more robust.
"""
if isinstance(value, torch.Tensor):
return float(value.reshape(-1)[0].detach().cpu().item())
if isinstance(value, np.ndarray):
return float(value.reshape(-1)[0])
if isinstance(value, (list, tuple)):
return float(np.asarray(value).reshape(-1)[0])
return float(value)
def _bool_from_env_value(value: Any) -> bool:
"""Normalize simulator boolean-like outputs into a plain Python bool."""
if isinstance(value, torch.Tensor):
return bool(value.reshape(-1)[0].detach().cpu().item())
if isinstance(value, np.ndarray):
return bool(value.reshape(-1)[0])
if isinstance(value, (list, tuple)):
return bool(np.asarray(value).reshape(-1)[0])
return bool(value)
def _make_output_dir(output_root: Path, scene: int) -> Path:
"""Create a timestamped output directory for one sim-eval run."""
timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
output_dir = output_root / f"scene{scene}_{timestamp}"
output_dir.mkdir(parents=True, exist_ok=True)
return output_dir
def main() -> None:
"""Entry point for one or more DROID simulation rollouts.
High-level flow:
1. parse command-line flags
2. bootstrap Isaac Lab / sim-evals imports
3. create the DROID environment
4. connect the DreamZero websocket client
5. run `episodes` rollouts
6. export videos + JSON summaries
"""
# Script-level arguments.
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--episodes", type=int, default=1, help="Number of episodes to run.")
parser.add_argument("--scene", type=int, default=1, help="DROID scene id (1/2/3).")
parser.add_argument("--host", type=str, default="127.0.0.1", help="vLLM DreamZero server host.")
parser.add_argument("--port", type=int, default=8000, help="vLLM DreamZero server port.")
parser.add_argument("--path", type=str, default=DEFAULT_PATH, help="Websocket path suffix.")
parser.add_argument(
"--open-loop-horizon",
type=int,
default=DEFAULT_OPEN_LOOP_HORIZON,
help="How many actions to consume locally before requesting the next chunk.",
)
parser.add_argument(
"--output-root",
type=Path,
default=DEFAULT_OUTPUT_ROOT,
help="Directory where videos and trajectory logs are stored.",
)
try:
from isaaclab.app import AppLauncher
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError(
"Optional DROID sim-eval client requires Isaac Lab (`isaaclab`). "
"Launch it from an Isaac Lab environment, e.g. via `isaaclab.sh -p`."
) from exc
# Let Isaac Lab inject its own runtime flags (e.g. headless, device).
AppLauncher.add_app_launcher_args(parser)
args = parser.parse_args()
# DreamZero sim-eval always needs camera observations enabled.
args.enable_cameras = True
# Boot Isaac Sim / Isaac Lab.
app_launcher = AppLauncher(args)
simulation_app = app_launcher.app
# Set defaults so the `finally` block can clean up safely even if an
# earlier step fails.
env = None
client = None
# Import simulator modules only *after* the app is launched. This matches
# Isaac Lab's required import ordering.
try:
import gymnasium as gym
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `gymnasium`.") from exc
try:
import sim_evals.environments # noqa: F401
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError(
"Optional DROID sim-eval client requires the external `sim-evals` package or checkout to be importable."
) from exc
try:
from isaaclab_tasks.utils import parse_env_cfg
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("Optional DROID sim-eval client requires `isaaclab_tasks`.") from exc
# Resolve output location and scene prompt.
output_dir = _make_output_dir(args.output_root.expanduser().resolve(), args.scene)
instruction = _scene_instruction(args.scene)
# Build the DROID environment configuration from `sim-evals`.
env_cfg = parse_env_cfg(
"DROID",
device=args.device,
num_envs=1,
use_fabric=True,
)
# Select one of the pre-authored scenes/tasks.
env_cfg.set_scene(args.scene)
env = gym.make("DROID", cfg=env_cfg)
# Upstream sim-evals resets twice so materials / cameras are fully ready.
obs, _ = env.reset()
obs, _ = env.reset()
# Connect the websocket policy client.
client = DreamZeroJointPosClient(
remote_host=args.host,
remote_port=args.port,
path=args.path,
open_loop_horizon=args.open_loop_horizon,
)
# Aggregated per-run results.
all_episode_summaries: list[dict[str, Any]] = []
max_steps = int(env.env.max_episode_length)
logging.info("DreamZero metadata: %s", client.metadata())
logging.info("Scene %s prompt: %s", args.scene, instruction)
logging.info("Writing outputs to %s", output_dir)
try:
# No gradients are needed in inference-only rollout mode.
with torch.no_grad():
for episode_index in range(args.episodes):
# Per-episode collectors.
frames: list[np.ndarray] = []
step_records: list[StepRecord] = []
episode_start = time.perf_counter()
server_time_s = 0.0
final_reward = 0.0
terminated = False
truncated = False
for step_index in range(max_steps):
# Ask the policy for the next action. Internally this may or
# may not trigger a real server request depending on whether
# the local chunk cache has been exhausted.
logging.debug("Episode %d step %d: requesting action", episode_index, step_index)
result = client.infer(obs, instruction)
logging.debug(
"Episode %d step %d: got action (server_call=%s latency=%s)",
episode_index,
step_index,
result["used_server_call"],
result["chunk_latency_s"],
)
# Save one visualization frame per simulator step.
frames.append(result["viz"])
# Isaac Lab expects batched actions, hence `[None]`.
action_tensor = torch.tensor(result["action"], dtype=torch.float32)[None]
logging.debug("Episode %d step %d: stepping env", episode_index, step_index)
obs, reward, term, trunc, info = env.step(action_tensor)
logging.debug("Episode %d step %d: env.step returned", episode_index, step_index)
logging.debug("Episode %d step %d: parsing reward/flags", episode_index, step_index)
logging.debug(
"Episode %d step %d: raw types reward=%s term=%s trunc=%s",
episode_index,
step_index,
type(reward).__name__,
type(term).__name__,
type(trunc).__name__,
)
# Normalize environment outputs into plain Python scalars so
# the rest of the code does not depend on wrapper-specific types.
reward_value = _scalar_from_env_value(reward)
term_value = _bool_from_env_value(term)
trunc_value = _bool_from_env_value(trunc)
logging.debug(
"Episode %d step %d: parsed reward=%s term=%s trunc=%s",
episode_index,
step_index,
reward_value,
term_value,
trunc_value,
)
# Keep scene-object capture optional. It is useful for
# debugging / success heuristics, but the rollout should not
# fail if the environment does not expose object roots.
scene_objects = _capture_scene_objects(env)
# Accumulate total server-side time only on steps that
# triggered a fresh chunk inference.
if result["chunk_latency_s"] is not None:
server_time_s += float(result["chunk_latency_s"])
# Materialize one JSON-serializable trajectory record.
logging.debug("Episode %d step %d: appending trajectory", episode_index, step_index)
step_records.append(
StepRecord(
step_index=step_index,
used_server_call=bool(result["used_server_call"]),
chunk_latency_s=(
float(result["chunk_latency_s"]) if result["chunk_latency_s"] is not None else None
),
action=[float(x) for x in np.asarray(result["action"], dtype=np.float32).tolist()],
joint_position=[
float(x) for x in np.asarray(result["joint_position"], dtype=np.float32).tolist()
],
gripper_position=[
float(x) for x in np.asarray(result["gripper_position"], dtype=np.float32).tolist()
],
reward=reward_value,
terminated=term_value,
truncated=trunc_value,
scene_objects=scene_objects,
)
)
logging.debug("Episode %d step %d: trajectory appended", episode_index, step_index)
# Track final status for summary export.
final_reward = reward_value
terminated = term_value
truncated = trunc_value
if term_value or trunc_value:
# End the rollout early if the environment terminates.
break
# Episode-level timing and video export.
episode_wall_time_s = time.perf_counter() - episode_start
video_path = output_dir / f"episode_{episode_index:02d}.mp4"
logging.info("Episode %d: writing video to %s", episode_index, video_path)
mediapy.write_video(video_path, frames, fps=15)
# Reset the policy server between episodes.
logging.info("Episode %d: sending reset", episode_index)
reset_response = client.reset()
final_objects = step_records[-1].scene_objects if step_records else {}
success_report = _maybe_infer_success(args.scene, final_objects)
# Assemble the per-episode summary that is written to
# `episode_XX.json`.
episode_summary = {
"episode_index": episode_index,
"prompt": instruction,
"video_path": str(video_path),
"steps_executed": len(step_records),
"max_steps": max_steps,
"terminated": terminated,
"truncated": truncated,
"final_reward": final_reward,
"server_calls": client.server_calls,
"server_time_s": server_time_s,
"episode_wall_time_s": episode_wall_time_s,
"avg_server_time_per_call_s": (
server_time_s / client.server_calls if client.server_calls else None
),
"reset_response": reset_response,
"success_report": success_report,
"server_metadata": client.metadata(),
"trajectory": [record.__dict__ for record in step_records],
}
_dump_json(output_dir / f"episode_{episode_index:02d}.json", episode_summary)
all_episode_summaries.append(episode_summary)
logging.info(
"Episode %d done: steps=%d wall=%.2fs server_calls=%d heuristic_success=%s",
episode_index,
len(step_records),
episode_wall_time_s,
client.server_calls,
success_report.get("heuristic_success"),
)
# Reset per-episode counters while keeping the client alive.
client.server_calls = 0
# Top-level run summary across all episodes.
summary = {
"scene": args.scene,
"prompt": instruction,
"episodes": args.episodes,
"host": args.host,
"port": args.port,
"path": args.path,
"device": args.device,
"output_dir": str(output_dir),
"server_metadata": client.metadata(),
"episodes_summary": all_episode_summaries,
}
_dump_json(output_dir / "summary.json", summary)
# Also print the summary to stdout so the caller can capture it in logs.
print(json.dumps(summary, ensure_ascii=False, indent=2))
finally:
# Best-effort cleanup. Avoid masking the main error if cleanup fails.
try:
if client is not None:
client.client.close()
except Exception:
pass
if env is not None:
env.close()
simulation_app.close()
if __name__ == "__main__":
# Keep the script-level logs readable. Per-step rollout details are still
# available via `DEBUG` if needed, but websocket / asyncio internals are
# usually too noisy for normal usage.
logging.basicConfig(level=logging.INFO)
logging.getLogger("websockets").setLevel(logging.WARNING)
logging.getLogger("asyncio").setLevel(logging.WARNING)
main()
@@ -1,86 +0,0 @@
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
os.environ.setdefault("MUJOCO_GL", "egl")
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
# Import base configs at module top level so the subclasses below are pickle-
# resolvable (worker processes import this module fresh via __main__).
from molmo_spaces.configs.policy_configs_baselines import ( # noqa: E402
DreamZeroPolicyConfig,
)
from molmo_spaces.evaluation.configs.evaluation_configs import ( # noqa: E402
DreamZeroPolicyEvalConfig,
)
# We only need to change the backend host and port to the vllm-host!
class DreamZeroVllmOmniPolicyConfig(DreamZeroPolicyConfig):
remote_config: dict = dict(host="127.0.0.1", port=8000)
class DreamZeroVllmOmniEvalConfig(DreamZeroPolicyEvalConfig):
policy_config: DreamZeroVllmOmniPolicyConfig = DreamZeroVllmOmniPolicyConfig()
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument(
"--benchmark_dir",
required=True,
help=(
"Path to a MolmoSpaces benchmark directory, for example "
"$MOLMOSPACES_BENCHMARK_DIR/20260327/ithor/FrankaCloseHardBench/"
"FrankaCloseHardBench_20260206_json_benchmark"
),
)
parser.add_argument("--max_episodes", type=int, default=1)
parser.add_argument("--task_horizon_steps", type=int, default=80)
parser.add_argument(
"--output_dir",
required=True,
help="Directory to write evaluation outputs (created if missing).",
)
parser.add_argument("--episode_idx", type=int, default=None)
args = parser.parse_args()
policy_config = DreamZeroVllmOmniPolicyConfig(remote_config=dict(host=args.host, port=args.port))
DreamZeroVllmOmniPolicyConfig.model_fields["remote_config"].default = policy_config.remote_config
DreamZeroVllmOmniEvalConfig.model_fields["policy_config"].default = policy_config
# Import after env vars are set so MuJoCo picks EGL.
from molmo_spaces.evaluation import run_evaluation
cfg_cls = DreamZeroVllmOmniEvalConfig
output_dir = args.output_dir
Path(output_dir).mkdir(parents=True, exist_ok=True)
print(f"[eval] benchmark_dir={args.benchmark_dir}")
print(f"[eval] max_episodes={args.max_episodes} task_horizon_steps={args.task_horizon_steps}")
print(f"[eval] remote policy: ws://{args.host}:{args.port}/v1/realtime/robot/openpi")
results = run_evaluation(
eval_config_cls=cfg_cls,
benchmark_dir=Path(args.benchmark_dir),
max_episodes=args.max_episodes,
task_horizon_steps=args.task_horizon_steps,
num_workers=1,
use_wandb=False,
output_dir=output_dir,
episode_idx=args.episode_idx,
)
print(f"[eval] success={results.success_count}/{results.total_count} ({results.success_rate:.1%})")
print(f"[eval] output_dir={results.output_dir}")
return 0
if __name__ == "__main__":
sys.exit(main())
@@ -1,311 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import argparse
import json
import logging
import sys
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
try:
import cv2
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("DreamZero OpenPI example requires `opencv-python`.") from exc
try:
example_dir = str(Path(__file__).resolve().parent)
removed_path = False
if sys.path and sys.path[0] == example_dir:
sys.path.pop(0)
removed_path = True
try:
from openpi_client.websocket_client_policy import WebsocketClientPolicy
finally:
if removed_path:
sys.path.insert(0, example_dir)
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("DreamZero OpenPI example requires `openpi-client`.") from exc
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 8000
DEFAULT_PATH = "/v1/realtime/robot/openpi"
DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
ACTION_HORIZON = 24
DEFAULT_ACTION_DIM = 8
RELATIVE_OFFSETS = [-23, -16, -8, 0]
REPO_ROOT = Path(__file__).resolve().parents[3]
ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
DEFAULT_VIDEO_DIR = REPO_ROOT / "outputs" / "dreamzero" / "assets"
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 _as_action_array(response: Any) -> np.ndarray:
if isinstance(response, dict) and response.get("type") == "error":
message = response.get("message", response)
raise RuntimeError(f"Inference failed: {message}")
return np.asarray(response, dtype=np.float32)
@dataclass(frozen=True)
class DreamZeroServerMetadata:
image_resolution: tuple[int, int]
n_external_cameras: int
needs_wrist_camera: bool
needs_stereo_camera: bool
needs_session_id: bool
action_space: str
@classmethod
def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
required_keys = (
"image_resolution",
"n_external_cameras",
"needs_wrist_camera",
"needs_stereo_camera",
"needs_session_id",
"action_space",
)
missing_keys = [key for key in required_keys if key not in payload]
if missing_keys:
raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
image_resolution = payload["image_resolution"]
if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
return cls(
image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
n_external_cameras=int(payload["n_external_cameras"]),
needs_wrist_camera=bool(payload["needs_wrist_camera"]),
needs_stereo_camera=bool(payload["needs_stereo_camera"]),
needs_session_id=bool(payload["needs_session_id"]),
action_space=str(payload["action_space"]),
)
def _openpi_uri(host: str, port: int, path: str) -> str:
if host.startswith(("ws://", "wss://")):
return host
return f"ws://{host}:{port}{path}"
def _close_policy_client(client: WebsocketClientPolicy) -> None:
ws = getattr(client, "_ws", None)
if ws is not None:
ws.close()
def load_all_frames(video_path: Path) -> np.ndarray:
cap = cv2.VideoCapture(str(video_path))
frames = []
while True:
ok, frame = cap.read()
if not ok:
break
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
cap.release()
if not frames:
raise RuntimeError(f"No frames loaded from {video_path}")
return np.stack(frames, axis=0)
def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
camera_frames: dict[str, np.ndarray] = {}
for camera_key, file_name in CAMERA_FILES.items():
video_path = video_dir / file_name
if not video_path.exists():
raise FileNotFoundError(
f"Missing DreamZero example asset: {video_path}. "
"Download the example videos with: "
f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
)
camera_frames[camera_key] = load_all_frames(video_path)
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],
*,
prompt: str,
session_id: str,
) -> dict[str, Any]:
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_demo_observations(
camera_frames: dict[str, np.ndarray],
*,
prompt: str,
session_id: str,
num_chunks: int = 2,
) -> list[dict[str, Any]]:
if num_chunks < 1:
raise ValueError("num_chunks must be at least 1")
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,
)
]
for indices in build_frame_schedule(total_frames, num_chunks - 1):
observations.append(
make_obs_from_video(
camera_frames,
indices,
prompt=prompt,
session_id=session_id,
)
)
return observations
def validate_session_result(
result: dict[str, Any],
*,
expected_action_horizon: int = ACTION_HORIZON,
expected_action_dim: int = DEFAULT_ACTION_DIM,
expected_action_count: int | None = None,
) -> None:
metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
if metadata.image_resolution != (180, 320):
raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
if metadata.n_external_cameras != 2:
raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
if not metadata.needs_wrist_camera:
raise AssertionError("DreamZero example expects wrist camera metadata")
if metadata.action_space != "joint_position":
raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
actions = result["actions"]
if expected_action_count is not None and len(actions) != expected_action_count:
raise AssertionError(f"Expected {expected_action_count} action tensors, got {len(actions)}")
if not actions:
raise AssertionError("Expected at least one action tensor")
for index, action in enumerate(actions):
if action.shape != (expected_action_horizon, expected_action_dim):
raise AssertionError(
f"Action {index} shape mismatch: expected "
f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
)
if not np.isfinite(action).all():
raise AssertionError(f"Action {index} contains non-finite values")
def run_policy_session(
*,
host: str = DEFAULT_HOST,
port: int = DEFAULT_PORT,
path: str = DEFAULT_PATH,
video_dir: Path = DEFAULT_VIDEO_DIR,
prompt: str = DEFAULT_PROMPT,
session_id: str | None = None,
num_chunks: int = 2,
) -> dict[str, Any]:
session_id = session_id or str(uuid.uuid4())
camera_frames = load_camera_frames(video_dir)
observations = build_demo_observations(
camera_frames,
prompt=prompt,
session_id=session_id,
num_chunks=num_chunks,
)
uri = _openpi_uri(host, port, path)
logging.info("Connecting to %s", uri)
client = WebsocketClientPolicy(host=uri)
try:
metadata = dict(client.get_server_metadata())
actions = [_as_action_array(client.infer(obs)) for obs in observations]
return {
"metadata": metadata,
"actions": actions,
"session_id": session_id,
}
finally:
_close_policy_client(client)
def format_action_summary(index: int, action: np.ndarray) -> str:
return (
f"Action {index}: shape={tuple(action.shape)} dtype={action.dtype} "
f"min={action.min():.6f} max={action.max():.6f}"
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="DreamZero OpenPI client example with downloaded real videos.")
parser.add_argument("--host", default=DEFAULT_HOST)
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
parser.add_argument("--path", default=DEFAULT_PATH)
parser.add_argument(
"--video-dir", type=Path, default=DEFAULT_VIDEO_DIR, help="Directory containing the three camera MP4 files."
)
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
parser.add_argument("--session-id", default=None)
parser.add_argument("--num-chunks", type=int, default=2)
return parser.parse_args()
def main() -> int:
args = parse_args()
logging.basicConfig(level=logging.INFO)
result = run_policy_session(
host=args.host,
port=args.port,
path=args.path,
video_dir=args.video_dir,
prompt=args.prompt,
session_id=args.session_id,
num_chunks=args.num_chunks,
)
validate_session_result(result, expected_action_count=args.num_chunks)
print("Server metadata:", json.dumps(result["metadata"], sort_keys=True))
for index, action in enumerate(result["actions"]):
print(format_action_summary(index, action))
print("Session ID:", result["session_id"])
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -1,24 +0,0 @@
#!/usr/bin/env bash
set -euo pipefail
MODEL="${MODEL:-GEAR-Dreams/DreamZero-DROID}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8000}"
DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/dreamzero_tp1_cfg2.yaml}"
SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-dreamzero-droid}"
args=(
serve
"$MODEL"
--omni
--host "$HOST"
--port "$PORT"
--served-model-name "$SERVED_MODEL_NAME"
--deploy-config "$DEPLOY_CONFIG"
--enforce-eager
--disable-log-stats
)
ATTENTION_BACKEND="${ATTENTION_BACKEND:-torch}" \
DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-TORCH_SDPA}" \
vllm "${args[@]}"
+5 -5
View File
@@ -22,14 +22,13 @@ export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
export HF_MODULES_CACHE=<transformers_modules_path>
```
Run from repository root. The 3-stage pipeline is resolved automatically from
the model type (`vllm_omni/deploy/dynin_omni.yaml`), so no config path is
required. Pass `--deploy-config` only to override the defaults:
Run from repository root:
```bash
vllm-omni serve snu-aidas/Dynin-Omni \
--omni \
--port 8091
--port 8091 \
--stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
```
If `vllm-omni` is not in PATH, run:
@@ -37,7 +36,8 @@ If `vllm-omni` is not in PATH, run:
```bash
PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
--omni \
--port 8091
--port 8091 \
--stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
```
Wait until the server logs show both `All stages initialized successfully` and
@@ -12,8 +12,7 @@ from typing import Any
import requests
from openai import OpenAI
from vllm.assets.audio import AudioAsset
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
from vllm.utils.argparse_utils import FlexibleArgumentParser
# Modify OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
@@ -494,7 +493,7 @@ def run_multimodal_generation(args) -> None:
def parse_args():
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--query-type",
"-q",
@@ -1,65 +0,0 @@
# vLLM-Omni · MiniCPM-o 4.5 Online Demo
Gradio-based web UI for **MiniCPM-o 4.5** served via `vllm-omni`'s
OpenAI-compatible endpoints.
The UI supports:
- **Inputs**: text prompt + optional image, audio (file or mic), video.
- **Outputs**: text + speech (WAV player).
## 1. Start the backend server
The deploy config auto-loads via `--omni`; the default
`vllm_omni/deploy/minicpmo_4_5.yaml` targets a 2-GPU layout (thinker on
GPU 0, talker + t2w sharing GPU 1). For other hardware layouts pick
one of the deploy variants below.
| deploy config | GPUs | Notes |
|---|---|---|
| `minicpmo_4_5.yaml` (default) | 2 | Thinker on GPU0, talker+t2w on GPU1. |
| `minicpmo_4_5_3gpu.yaml` | 3 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. |
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
Default (2-GPU):
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--trust-remote-code \
--host 0.0.0.0 --port 8099
```
Other layouts via `--deploy-config`:
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \
--trust-remote-code \
--host 0.0.0.0 --port 8099
```
## 2. Launch the Gradio demo
```bash
bash examples/online_serving/minicpmo/run_gradio_demo.sh
# Or run the python entry point directly:
python examples/online_serving/minicpmo/gradio_demo.py \
--minicpmo45-api-base http://localhost:8099/v1 \
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
--port 7862
```
Open `http://<host>:7862` in a browser.
## Notes
- **TTS trigger**: the demo sets
`extra_body.chat_template_kwargs.use_tts_template=True`, which appends
`<|tts_bos|>` to the assistant prefix.
- Uncheck **"Generate speech output (TTS)"** to get text-only responses
(faster).
- The audio output is the raw WAV returned by the stage-1 talker +
Token2Wav; sample rate is 24 kHz.
- Video input is forwarded as a base64 `video_url` entry; the server needs
decord/torchvision to decode it.
@@ -1,580 +0,0 @@
"""Gradio demo for MiniCPM-o 4.5 online serving through vllm-omni.
Supports:
- Text / Image / Audio / Video inputs (OpenAI chat-completions multimodal format).
- Text + audio outputs, returned side-by-side in the UI.
Usage:
# run a 4.5 server on :8099
python gradio_demo.py \
--minicpmo45-api-base http://localhost:8099/v1 \
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
--port 7862
"""
from __future__ import annotations
import argparse
import base64
import io
import os
from collections.abc import Iterable
from pathlib import Path
from typing import Any
import gradio as gr
import numpy as np
import soundfile as sf
from openai import OpenAI
from PIL import Image
MINICPMO45 = "MiniCPM-o-4.5"
# Default reference audio per model, relative to the model path. Used to build
# the audio_assistant system prompt so that the LLM adopts the "speak in this
# voice" persona (matches the official MiniCPM-o demo).
_DEFAULT_REF_AUDIO = {
MINICPMO45: "assets/HT_ref_audio.wav",
}
# audio_assistant prompt text, mirrors MiniCPMO.get_sys_prompt("audio_assistant")
# in the MiniCPM-o-Demo reference implementation.
_AUDIO_ASSISTANT_PROMPT = {
"zh": {
"prefix": "模仿音频样本的音色并生成新的内容。",
"suffix": (
"你的任务是用这种声音模式来当一个助手。请认真、高质量地回复用户的问题。"
"请用高自然度的方式和用户聊天。你是由面壁智能开发的人工智能助手:面壁小钢炮。"
),
},
"en": {
"prefix": "Use the voice in the audio prompt to synthesize new content.",
"suffix": "You are a helpful assistant with the above voice style.",
},
}
# ---------------------------------------------------------------------------
# Media helpers
# ---------------------------------------------------------------------------
def image_to_base64_data_url(image: Image.Image) -> str:
buf = io.BytesIO()
if image.mode != "RGB":
image = image.convert("RGB")
image.save(buf, format="JPEG", quality=92)
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:image/jpeg;base64,{b64}"
def audio_to_base64_data_url(audio_np: np.ndarray, sample_rate: int) -> str:
if audio_np.dtype != np.int16:
if audio_np.dtype in (np.float32, np.float64):
audio_np = np.clip(audio_np, -1.0, 1.0)
audio_np = (audio_np * 32767).astype(np.int16)
else:
audio_np = audio_np.astype(np.int16)
buf = io.BytesIO()
sf.write(buf, audio_np, sample_rate, format="WAV")
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:audio/wav;base64,{b64}"
def ref_audio_to_data_url(path: str, target_sr: int = 16000, max_s: float | None = 8.0) -> str:
"""Load a file, downmix to mono, resample to 16kHz, optionally truncate, and
return a data-URL usable as audio_url content in the chat API.
MiniCPM-o expects 16kHz mono audio input; long reference audios can waste
tokens so we cap at max_s seconds by default (3-10s is the recommended range).
"""
data, sr = sf.read(path)
if data.ndim > 1:
data = data[:, 0]
data = data.astype(np.float32)
if sr != target_sr:
try:
# Demo client only; librosa here resamples audio before sending to
# the server. The fallback below covers environments without it.
import librosa # noqa: TID251
data = librosa.resample(data, orig_sr=sr, target_sr=target_sr)
except ImportError:
from math import gcd
g = gcd(sr, target_sr)
up, down = target_sr // g, sr // g
idx = (np.arange(len(data) * up) / up).astype(np.int64)
data = data[idx][::down]
sr = target_sr
if max_s is not None and len(data) > int(max_s * sr):
data = data[: int(max_s * sr)]
return audio_to_base64_data_url(data, sr)
def video_to_base64_data_url(path: str) -> str:
p = Path(path)
if not p.exists():
raise FileNotFoundError(f"Video not found: {path}")
ext = p.suffix.lower()
mime = {
".mp4": "video/mp4",
".webm": "video/webm",
".mov": "video/quicktime",
".avi": "video/x-msvideo",
".mkv": "video/x-matroska",
}.get(ext, "video/mp4")
b64 = base64.b64encode(p.read_bytes()).decode("utf-8")
return f"data:{mime};base64,{b64}"
def process_audio_input(audio_input: Any | None) -> tuple[np.ndarray, int] | None:
"""Normalize Gradio audio input to (np.ndarray mono float32, sample_rate).
Gradio's `gr.Audio(type="numpy")` yields (sr, ndarray); we also accept
filesystem paths and a few nested tuple variants that appear in different
Gradio versions.
"""
if audio_input is None:
return None
def _from_path(p: str) -> tuple[np.ndarray, int] | None:
if not p:
return None
fp = Path(p)
if not fp.exists():
return None
data, sr = sf.read(fp)
if data.ndim > 1:
data = data[:, 0]
return data.astype(np.float32), int(sr)
audio_np: np.ndarray | None = None
sr: int | None = None
if isinstance(audio_input, tuple) and len(audio_input) == 2:
a, b = audio_input
if isinstance(a, int | float) and isinstance(b, np.ndarray):
sr, audio_np = int(a), b
elif isinstance(a, str):
loaded = _from_path(a)
if loaded is not None:
audio_np, sr = loaded
elif isinstance(audio_input, str):
loaded = _from_path(audio_input)
if loaded is not None:
audio_np, sr = loaded
if audio_np is None or sr is None:
return None
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
return audio_np.astype(np.float32), sr
# ---------------------------------------------------------------------------
# Per-model inference config
# ---------------------------------------------------------------------------
class ModelEndpoint:
"""Endpoint + TTS-trigger convention for one MiniCPM-o variant."""
def __init__(self, name: str, api_base: str, model_path: str):
self.name = name
self.api_base = api_base
self.model_path = model_path
self.client = OpenAI(api_key="EMPTY", base_url=api_base)
self._ref_audio_data_url: str | None = None
@property
def ref_audio_path(self) -> str | None:
rel = _DEFAULT_REF_AUDIO.get(self.name)
if not rel:
return None
full = os.path.join(self.model_path, rel)
return full if os.path.exists(full) else None
def get_ref_audio_data_url(self) -> str | None:
"""Lazily load and cache the default reference audio as a 16kHz mono
data URL. Returns None if the reference audio file is missing."""
if self._ref_audio_data_url is not None:
return self._ref_audio_data_url
p = self.ref_audio_path
if not p:
print(f"[{self.name}] no default reference audio found; falling back to text-only system prompt")
return None
try:
self._ref_audio_data_url = ref_audio_to_data_url(p)
print(f"[{self.name}] loaded reference audio: {p}")
except Exception as e:
print(f"[{self.name}] failed to load reference audio {p}: {e}")
return None
return self._ref_audio_data_url
def build_extras(self) -> dict[str, Any]:
"""Extra kwargs passed to chat.completions.create for this model.
4.5: use_tts_template=True chat-template flag appends <|tts_bos|>.
"""
if self.name == MINICPMO45:
return {
"extra_body": {
"chat_template_kwargs": {"use_tts_template": True},
"modalities": ["text", "audio"],
},
}
return {}
# ---------------------------------------------------------------------------
# System prompt (matches MiniCPMO.get_sys_prompt("audio_assistant") in the
# official MiniCPM-o-Demo). When TTS is disabled we fall back to a plain
# text-only system prompt.
# ---------------------------------------------------------------------------
def build_audio_assistant_system(
ref_audio_data_url: str | None,
language: str = "zh",
) -> dict[str, Any]:
"""Return an OpenAI-style system message carrying the reference audio."""
txt = _AUDIO_ASSISTANT_PROMPT.get(language, _AUDIO_ASSISTANT_PROMPT["zh"])
if ref_audio_data_url is None:
return {
"role": "system",
"content": (
"You are MiniCPM-o, a helpful multimodal assistant that can "
"understand images, audio and video, and respond in text and speech."
),
}
return {
"role": "system",
"content": [
{"type": "text", "text": txt["prefix"]},
{"type": "audio_url", "audio_url": {"url": ref_audio_data_url}},
{"type": "text", "text": txt["suffix"]},
],
}
# ---------------------------------------------------------------------------
# Core inference function
# ---------------------------------------------------------------------------
def run_inference(
endpoint: ModelEndpoint,
user_prompt: str,
audio_tuple: tuple[np.ndarray, int] | None,
image_pil: Image.Image | None,
video_path: str | None,
enable_tts: bool,
max_tokens: int,
temperature: float,
) -> Iterable[tuple[str, tuple[int, np.ndarray] | None]]:
user_prompt = user_prompt or ""
if not user_prompt.strip() and not audio_tuple and not image_pil and not video_path:
yield "Please provide a text prompt or at least one multimodal input.", None
return
try:
content: list[dict[str, Any]] = []
if audio_tuple is not None:
audio_np, sr = audio_tuple
content.append(
{"type": "audio_url", "audio_url": {"url": audio_to_base64_data_url(audio_np, sr)}},
)
if image_pil is not None:
img = image_pil if image_pil.mode == "RGB" else image_pil.convert("RGB")
content.append(
{"type": "image_url", "image_url": {"url": image_to_base64_data_url(img)}},
)
if video_path:
content.append(
{"type": "video_url", "video_url": {"url": video_to_base64_data_url(video_path)}},
)
if user_prompt.strip():
content.append({"type": "text", "text": user_prompt})
if enable_tts:
ref_url = endpoint.get_ref_audio_data_url()
system_msg = build_audio_assistant_system(ref_url, language="zh")
else:
system_msg = {
"role": "system",
"content": (
"You are MiniCPM-o, a helpful multimodal assistant that can "
"understand images, audio and video, and respond in text and speech."
),
}
messages = [
system_msg,
{"role": "user", "content": content},
]
kwargs: dict[str, Any] = {
"model": endpoint.model_path,
"messages": messages,
"max_tokens": int(max_tokens),
"temperature": float(temperature),
}
if enable_tts:
extras = endpoint.build_extras()
# Merge extra_body dicts cleanly.
extra_body = dict(extras.get("extra_body", {}))
kwargs["extra_body"] = extra_body
else:
# Text-only: suppress audio output.
kwargs["extra_body"] = {"modalities": ["text"]}
completion = endpoint.client.chat.completions.create(**kwargs)
text_parts: list[str] = []
audio_out: tuple[int, np.ndarray] | None = None
for choice in completion.choices:
msg = choice.message
if getattr(msg, "content", None):
text_parts.append(msg.content)
audio_obj = getattr(msg, "audio", None)
if audio_obj and getattr(audio_obj, "data", None):
raw = base64.b64decode(audio_obj.data)
if len(raw) > 80:
data, sr = sf.read(io.BytesIO(raw))
if data.ndim > 1:
data = data[:, 0]
audio_out = (int(sr), data.astype(np.float32))
text_response = "\n\n".join(t for t in text_parts if t) or "(no text returned)"
yield text_response, audio_out
except Exception as exc: # noqa: BLE001
yield f"Inference failed: {type(exc).__name__}: {exc}", None
# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
DEFAULT_PROMPTS = [
"你好,请用一句话介绍北京。",
"Describe what you see in the image.",
"Transcribe the audio and translate it into English.",
"What is happening in this video?",
]
def build_interface(endpoints: dict[str, ModelEndpoint], default_model: str) -> gr.Blocks:
model_choices = list(endpoints.keys())
css = """
#generate-btn button { width: 100%; font-size: 1.05rem; }
.media-row > div { min-height: 260px; }
.model-banner { color: #555; font-size: 0.9rem; }
"""
with gr.Blocks(css=css, title="vLLM-Omni · MiniCPM-o Demo") as demo:
gr.Markdown(
"# vLLM-Omni · MiniCPM-o Online Demo\n"
"Multimodal chat (text / image / audio / video in; text + speech out) for "
"MiniCPM-o 4.5 served by vllm-omni."
)
with gr.Row():
model_dd = gr.Dropdown(
label="Model",
choices=model_choices,
value=default_model,
scale=2,
)
endpoint_info = gr.Markdown(
elem_classes="model-banner",
value=f"`{endpoints[default_model].api_base}` → `{endpoints[default_model].model_path}`",
)
with gr.Row():
prompt_box = gr.Textbox(
label="Text Prompt",
placeholder="Ask the model anything. Combine with image / audio / video below.",
lines=3,
value=DEFAULT_PROMPTS[0],
scale=3,
)
with gr.Column(scale=1):
tts_checkbox = gr.Checkbox(label="Generate speech output (TTS)", value=True)
max_tokens_slider = gr.Slider(label="Max tokens", minimum=32, maximum=4096, value=1024, step=32)
temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=1.5, value=0.7, step=0.05)
with gr.Row(elem_classes="media-row"):
image_input = gr.Image(label="Image (optional)", type="pil", sources=["upload"])
audio_input = gr.Audio(
label="Audio (optional)",
type="numpy",
sources=["upload", "microphone"],
)
video_input = gr.Video(label="Video (optional)", sources=["upload"])
with gr.Row():
generate_btn = gr.Button("Generate", variant="primary", elem_id="generate-btn")
clear_btn = gr.Button("Clear")
with gr.Row():
text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)
gr.Examples(
examples=[
["你好,请用一句话介绍北京。"],
["Describe what you see in detail."],
["请先识别这段语音的内容,再用中文回复说话人。"],
["What is happening in this video?"],
],
inputs=[prompt_box],
label="Quick prompts",
)
def _on_model_change(name: str) -> str:
ep = endpoints[name]
return f"`{ep.api_base}` → `{ep.model_path}`"
model_dd.change(_on_model_change, inputs=[model_dd], outputs=[endpoint_info])
def _run(
model_name: str,
prompt: str,
image_pil: Image.Image | None,
audio_in: Any,
video_path: str | None,
enable_tts: bool,
max_tokens: int,
temperature: float,
):
ep = endpoints[model_name]
audio_tuple = process_audio_input(audio_in)
yield from run_inference(
ep,
prompt,
audio_tuple,
image_pil,
video_path,
enable_tts,
max_tokens,
temperature,
)
generate_btn.click(
_run,
inputs=[
model_dd,
prompt_box,
image_input,
audio_input,
video_input,
tts_checkbox,
max_tokens_slider,
temperature_slider,
],
outputs=[text_output, audio_output],
)
def _clear():
return "", None, None, None, "", None
clear_btn.click(
_clear,
outputs=[
prompt_box,
image_input,
audio_input,
video_input,
text_output,
audio_output,
],
)
demo.queue()
return demo
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--minicpmo45-api-base", default=os.environ.get("MINICPMO45_API_BASE", "http://localhost:8099/v1"))
p.add_argument(
"--minicpmo45-model",
default=os.environ.get("MINICPMO45_MODEL", "openbmb/MiniCPM-o-4_5"),
)
p.add_argument("--host", default="0.0.0.0")
p.add_argument("--port", type=int, default=7862)
p.add_argument("--share", action="store_true")
p.add_argument(
"--ssl-certfile",
default=os.environ.get("GRADIO_SSL_CERTFILE", ""),
help="Path to TLS certificate (PEM). Browsers require HTTPS to access the microphone.",
)
p.add_argument(
"--ssl-keyfile",
default=os.environ.get("GRADIO_SSL_KEYFILE", ""),
help="Path to TLS private key (PEM).",
)
p.add_argument(
"--ssl-verify",
action="store_true",
default=os.environ.get("GRADIO_SSL_VERIFY", "0") == "1",
help="Require browser to verify the certificate (off by default; self-signed certs should stay off).",
)
return p.parse_args()
def _ping(api_base: str, timeout: float = 3.0) -> bool:
import urllib.request
try:
with urllib.request.urlopen(api_base.rstrip("/").replace("/v1", "/health"), timeout=timeout):
return True
except Exception:
return False
def main() -> None:
args = parse_args()
endpoints: dict[str, ModelEndpoint] = {}
if args.minicpmo45_api_base and args.minicpmo45_model:
ok = _ping(args.minicpmo45_api_base)
status = "reachable" if ok else "not reachable (will error on generate)"
print(f"[4.5] {args.minicpmo45_api_base} ({status})")
endpoints[MINICPMO45] = ModelEndpoint(MINICPMO45, args.minicpmo45_api_base, args.minicpmo45_model)
if not endpoints:
raise SystemExit("No endpoints configured. Pass --minicpmo45-api-base/--minicpmo45-model.")
default_model = next(iter(endpoints.keys()))
demo = build_interface(endpoints, default_model)
launch_kwargs: dict = {
"server_name": args.host,
"server_port": args.port,
"share": args.share,
}
if args.ssl_certfile and args.ssl_keyfile:
launch_kwargs.update(
ssl_certfile=args.ssl_certfile,
ssl_keyfile=args.ssl_keyfile,
ssl_verify=args.ssl_verify,
)
print(f"[tls] cert={args.ssl_certfile} key={args.ssl_keyfile} verify={args.ssl_verify}")
demo.launch(**launch_kwargs)
if __name__ == "__main__":
main()
@@ -1,36 +0,0 @@
#!/bin/bash
# Launch the MiniCPM-o 4.5 gradio demo.
#
# Prereq:
# Start a vllm-omni OpenAI server for MiniCPM-o 4.5 on :8099 (see the
# 8x4090 stage config under vllm_omni/model_executor/stage_configs).
set -e
HERE="$(cd "$(dirname "$0")" && pwd)"
: "${MINICPMO45_API_BASE:=http://localhost:8099/v1}"
: "${MINICPMO45_MODEL:=openbmb/MiniCPM-o-4_5}"
: "${GRADIO_HOST:=0.0.0.0}"
: "${GRADIO_PORT:=7862}"
# HTTPS (browsers require a secure context for microphone access).
# Set GRADIO_SSL_CERTFILE / GRADIO_SSL_KEYFILE to enable TLS.
: "${GRADIO_SSL_CERTFILE:=}"
: "${GRADIO_SSL_KEYFILE:=}"
export MINICPMO45_API_BASE MINICPMO45_MODEL
SSL_ARGS=()
if [ -n "$GRADIO_SSL_CERTFILE" ] && [ -n "$GRADIO_SSL_KEYFILE" ] \
&& [ -f "$GRADIO_SSL_CERTFILE" ] && [ -f "$GRADIO_SSL_KEYFILE" ]; then
SSL_ARGS=(--ssl-certfile "$GRADIO_SSL_CERTFILE" --ssl-keyfile "$GRADIO_SSL_KEYFILE")
echo "HTTPS enabled: cert=$GRADIO_SSL_CERTFILE key=$GRADIO_SSL_KEYFILE"
else
echo "HTTPS disabled (cert/key not found). Microphone won't work on remote browsers."
fi
exec python "$HERE/gradio_demo.py" \
--minicpmo45-api-base "$MINICPMO45_API_BASE" \
--minicpmo45-model "$MINICPMO45_MODEL" \
--host "$GRADIO_HOST" \
--port "$GRADIO_PORT" \
"${SSL_ARGS[@]}"
Loaded 100 of 517 files, more files were not shown because too many files have changed in this diff. Show more