diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 96aa9e8ba..0b880ee77 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1,14 +1,14 @@ # Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci # from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly -# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly +# exception in upload_pipeline.py (nightly-test label or main NIGHTLY=1); otherwise upload-nightly # can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build. -# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml. +# Document 2: appended after `---`; same file, read by upload_pipeline.py (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml. steps: - label: ":github: Resolve skip-ci & upload pipeline" key: upload-ci-pipeline commands: - - "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh" + - "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml" agents: queue: "cpu_queue_premerge" @@ -31,7 +31,7 @@ steps: key: upload-ready-pipeline if: __UPLOAD_READY_IF__ commands: - - buildkite-agent pipeline upload .buildkite/test-ready.yml + - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml agents: queue: "cpu_queue_premerge" @@ -41,7 +41,7 @@ steps: key: upload-merge-pipeline if: __UPLOAD_MERGE_IF__ commands: - - buildkite-agent pipeline upload .buildkite/test-merge.yml + - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/rebase-pipeline.yaml b/.buildkite/rebase-pipeline.yaml index edd98fafd..f8f4aa217 100644 --- a/.buildkite/rebase-pipeline.yaml +++ b/.buildkite/rebase-pipeline.yaml @@ -14,7 +14,7 @@ steps: depends_on: image-build key: upload-ready-pipeline commands: - - buildkite-agent pipeline upload .buildkite/test-ready.yml + - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml agents: queue: "cpu_queue_premerge" @@ -23,7 +23,7 @@ steps: depends_on: image-build key: upload-merge-pipeline commands: - - buildkite-agent pipeline upload .buildkite/test-merge.yml + - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/scripts/upload_pipeline.py b/.buildkite/scripts/upload_pipeline.py new file mode 100644 index 000000000..1a1254dbc --- /dev/null +++ b/.buildkite/scripts/upload_pipeline.py @@ -0,0 +1,308 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Render and optionally upload Buildkite pipeline YAML with diff-aware logic. + +Bootstrap mode (pipeline.yml with __IMAGE_BUILD_IF__ placeholders): + - Detect docs-only PR/main changes and substitute skip-ci ``if`` expressions. + +Test pipeline mode (e.g. test-merge.yml): + - Drop steps or groups whose ``source_file_dependencies`` do not match changed files. + - ``source_file_dependencies`` may be set on a leaf step or on a group step. + - Always strip ``source_file_dependencies`` from uploaded YAML (Buildkite does not + recognize this uploader-only key). + +Usage: + python3 upload_pipeline.py [--upload] + + # 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 "") + + +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()) diff --git a/.buildkite/scripts/upload_pipeline_with_skip_ci.sh b/.buildkite/scripts/upload_pipeline_with_skip_ci.sh deleted file mode 100644 index c19fe7063..000000000 --- a/.buildkite/scripts/upload_pipeline_with_skip_ci.sh +++ /dev/null @@ -1,146 +0,0 @@ -#!/usr/bin/env bash -# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML -# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time. -# -# Usage: upload_pipeline_with_skip_ci.sh [pipeline_yaml] -# pipeline_yaml: path relative to repo root or absolute (default: .buildkite/pipeline.yml) -set -euo pipefail - -ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" -PIPELINE_ARG="${1:-.buildkite/pipeline.yml}" -if [[ "${PIPELINE_ARG}" = /* ]]; then - PIPELINE_YML="${PIPELINE_ARG}" -else - PIPELINE_YML="${ROOT}/${PIPELINE_ARG}" -fi - -# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr. -is_docs_only_change() { - local file_path - local has_any=0 - - while IFS= read -r file_path; do - [[ -z "${file_path}" ]] && continue - has_any=1 - - if [[ "${file_path}" == docs/* ]]; then - continue - fi - if [[ "${file_path}" == *.md ]]; then - continue - fi - if [[ "${file_path}" == "mkdocs.yml" ]]; then - continue - fi - return 1 - done - - [[ "${has_any}" -eq 1 ]] -} - -resolve_skip_ci() { - local is_pr_build=0 - local files - local base_branch base_ref - - if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then - is_pr_build=1 - fi - - if [[ "${is_pr_build}" -eq 1 ]]; then - base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" - if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then - echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2 - git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true - fi - - base_ref="" - if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then - base_ref="origin/${base_branch}" - elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then - base_ref="${base_branch}" - else - echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2 - echo -n 0 - return 0 - fi - - if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then - echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2 - echo -n 0 - return 0 - fi - elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then - if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then - echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2 - echo -n 0 - return 0 - fi - if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then - echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2 - echo -n 0 - return 0 - fi - else - echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2 - echo -n 0 - return 0 - fi - - if is_docs_only_change <<< "${files}"; then - echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2 - echo -n 1 - return 0 - fi - - echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2 - echo -n 0 -} - -SKIP_CI="$(resolve_skip_ci)" - -if [[ ! -f "${PIPELINE_YML}" ]]; then - echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2 - exit 1 -fi - -export ROOT SKIP_CI PIPELINE_YML -python3 <<'PY' | buildkite-agent pipeline upload -import os -import pathlib - -path = pathlib.Path(os.environ["PIPELINE_YML"]) -text = path.read_text(encoding="utf-8") -sep = "\n---\n" -# Two supported layouts: -# - multi-doc: doc 1 = bootstrap (loaded by Buildkite), doc 2 (after `---`) = uploaded steps. -# - single-doc: whole file is the uploaded steps; caller (UI init step) is the bootstrap. -if sep in text: - _, continuation = text.split(sep, 1) -else: - continuation = text - -skip = os.environ.get("SKIP_CI") == "1" -# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or -# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too. -nightly_only = ( - '(build.pull_request.labels includes "nightly-test") ' - '|| (build.branch == "main" && build.env("NIGHTLY") == "1")' -) -# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only. -if skip: - rep = f"'{nightly_only}'" - ready_rep = "'false'" - merge_rep = "'false'" -else: - rep = "'true'" - ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'" - merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\" && build.env(\"WEEKLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'" -rendered = ( - continuation - .replace("__IMAGE_BUILD_IF__", rep) - .replace("__UPLOAD_READY_IF__", ready_rep) - .replace("__UPLOAD_MERGE_IF__", merge_rep) -) -print(rendered, end="") -PY diff --git a/.buildkite/test-amd-merge.yml b/.buildkite/test-amd-merge.yml index b986d687c..1a5217b84 100644 --- a/.buildkite/test-amd-merge.yml +++ b/.buildkite/test-amd-merge.yml @@ -1,54 +1,111 @@ +env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + HF_HUB_DOWNLOAD_TIMEOUT: 300 + HF_HUB_ETAG_TIMEOUT: 60 + steps: -- label: "Simple · Diffusion & Model Executor Test" - agent_pool: mi325_1 - depends_on: amd-build - timeout_in_minutes: 50 - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_ROCM_USE_AITER=0 - # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM - # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" +- group: ":card_index_dividers: Simple Test" + steps: + - label: "Simple · Diffusion & Model Executor Test" + agent_pool: mi325_1 + depends_on: amd-build + timeout_in_minutes: 50 + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM + # that is not supported on CPU + - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" -- label: "Simple · Other Test" - agent_pool: mi325_1 - depends_on: amd-build - timeout_in_minutes: 50 - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_ROCM_USE_AITER=0 - # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM - # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - label: "Simple · Other Test" + agent_pool: mi325_1 + depends_on: amd-build + timeout_in_minutes: 50 + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM + # that is not supported on CPU + - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" -- label: "Diffusion Images API LoRA E2E" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py +- group: ":card_index_dividers: Diffusion Test" + steps: + - label: "Diffusion · LoRA Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py -- label: "Diffusion Model CPU offloading Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_LOGGING_LEVEL=DEBUG - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - | - timeout 20m bash -c ' - set +e - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py - EXIT1=\$? - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py - EXIT2=\$? - exit \$((EXIT1 | EXIT2)) - ' + - label: "Diffusion Model CPU offloading Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - | + timeout 20m bash -c ' + set +e + pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py + EXIT1=\$? + pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py + EXIT2=\$? + exit \$((EXIT1 | EXIT2)) + ' + + - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)" + agent_pool: mi325_4 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py + - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py + + # merge-only tests + - label: "Diffusion Tensor Parallelism Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py + + - label: "Diffusion Model Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py + + - label: "Diffusion GPU Worker Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py + + - label: "Diffusion · Wan22 Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - | + timeout 40m bash -c " + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' + " ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274 # - label: "Audio Generation Model Test" @@ -62,42 +119,6 @@ steps: # - export VLLM_WORKER_MULTIPROC_METHOD=spawn # - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model -- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)" - agent_pool: mi325_4 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_LOGGING_LEVEL=DEBUG - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py - - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py - -# merge-only tests -- label: "Diffusion Tensor Parallelism Test" - agent_pool: mi325_2 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py - -- label: "Diffusion Model Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py - -- label: "Diffusion GPU Worker Test" - agent_pool: mi325_2 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py - - label: "Engine Test" agent_pool: mi325_1 depends_on: amd-build @@ -111,50 +132,65 @@ steps: # This model seems to be important, so I am trying to run it here. # Since this is from test-nightly.yml, # I keep this non-blocking. -- label: "Omni Model Test Qwen2-5-Omni" - agent_pool: mi325_4 - depends_on: amd-build - mirror_hardwares: [amdproduction] - commands: - - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model" +- group: ":card_index_dividers: E2E Test" + steps: + - label: "Omni Model Test Qwen2-5-Omni" + agent_pool: mi325_4 + depends_on: amd-build + mirror_hardwares: [amdproduction] + commands: + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model" -- label: "Omni Model Test Qwen3-Omni" - agent_pool: mi325_2 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' + - label: "Omni Model Test Qwen3-Omni" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' -- label: "Qwen3-TTS CustomVoice E2E Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 20m bash -c ' - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" - ' + - label: "Qwen3-TTS CustomVoice E2E Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - | + timeout 20m bash -c ' + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" + ' -- label: "Qwen3-TTS Base E2E Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 30m bash -c ' - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" - ' + - label: "Qwen3-TTS Base E2E Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - | + timeout 30m bash -c ' + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" + ' + + - label: "Voxtral-TTS E2E Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - | + timeout 20m bash -c ' + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" + ' # TODO: Bagel test on ROCm is very unstable. @tjtanaa # Need to debug before reneable numerical changes across large PRs @@ -195,28 +231,3 @@ steps: # - export VLLM_WORKER_MULTIPROC_METHOD=spawn # - export VLLM_ROCM_USE_AITER_RMSNORM=0 # - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm" - -- label: "Diffusion · Wan22 Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 40m bash -c " - export VLLM_IMAGE_FETCH_TIMEOUT=60 - pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' - " - -- label: "Voxtral-TTS E2E Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 20m bash -c ' - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" - ' diff --git a/.buildkite/test-amd-ready.yaml b/.buildkite/test-amd-ready.yaml index 9951544bd..9da9507c8 100644 --- a/.buildkite/test-amd-ready.yaml +++ b/.buildkite/test-amd-ready.yaml @@ -1,42 +1,33 @@ +env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + HF_HUB_DOWNLOAD_TIMEOUT: 300 + HF_HUB_ETAG_TIMEOUT: 60 + steps: -- label: "Simple · Diffusion & Model Executor Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_ROCM_USE_AITER=0 - # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM - # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" +- group: ":card_index_dividers: Simple Test" + steps: + - label: "Simple · Diffusion & Model Executor Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM + # that is not supported on CPU + - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" -- label: "Simple · Other Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_ROCM_USE_AITER=0 - # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM - # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" - -- label: "Voxtral TTS CUDA Unit Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py - -- label: "Diffusion · Batch Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model" + - label: "Simple · Other Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM + # that is not supported on CPU + - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" - label: "Custom Pipeline Test" agent_pool: mi325_1 @@ -46,23 +37,92 @@ steps: commands: - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model" -- label: "Diffusion Model CPU offloading Test" + +- label: "Voxtral TTS CUDA Unit Test" agent_pool: mi325_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - - export VLLM_LOGGING_LEVEL=DEBUG - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - | - timeout 20m bash -c ' - set +e - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py - EXIT1=\$? - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py - EXIT2=\$? - exit \$((EXIT1 | EXIT2)) - ' + - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py + +- group: ":card_index_dividers: Diffusion Test" + steps: + - label: "Diffusion · Batch Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model" + + - label: "Diffusion · Offloader Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - | + timeout 20m bash -c ' + set +e + pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py + EXIT1=\$? + pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py + EXIT2=\$? + exit \$((EXIT1 | EXIT2)) + ' + + - label: "Diffusion · Cache Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4" + + - label: "Diffusion Sequence Parallelism Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model + + - label: "Diffusion GPU Worker Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py + + - label: "Diffusion · Model Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model" + + - label: "Diffusion · Wan22 Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - | + timeout 40m bash -c " + export VLLM_TEST_CLEAN_GPU_MEMORY=1 + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model' + " ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274 # - label: "Audio Generation Model Test" @@ -76,56 +136,6 @@ steps: # - export VLLM_WORKER_MULTIPROC_METHOD=spawn # - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model -- label: "Diffusion Cache Backend Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export GPU_ARCHS=gfx942 - - export VLLM_LOGGING_LEVEL=DEBUG - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4" - -- label: "Diffusion Sequence Parallelism Test" - agent_pool: mi325_2 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_LOGGING_LEVEL=DEBUG - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model - -- label: "Diffusion GPU Worker Test" - agent_pool: mi325_2 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py - -- label: "Diffusion · Model Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model" - -- label: "Diffusion · Wan22 Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 40m bash -c " - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - export VLLM_IMAGE_FETCH_TIMEOUT=60 - pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model' - " - - label: "Engine Test" agent_pool: mi325_1 depends_on: amd-build @@ -149,43 +159,56 @@ steps: # - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py # - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model" -- label: "MiMo-Audio E2E Test with H100" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 30m bash -c ' - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model" - ' +- group: ":card_index_dividers: E2E Test" + steps: + - label: "MiMo-Audio E2E Test with H100" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - | + timeout 30m bash -c ' + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model" + ' -- label: "Qwen3-TTS E2E Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export VLLM_LOGGING_LEVEL=DEBUG - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model" + - label: "Qwen3-TTS E2E Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model" -- label: "Voxtral-TTS E2E Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 20m bash -c ' - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" - pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" - ' + - label: "Voxtral-TTS E2E Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - | + timeout 20m bash -c ' + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" + pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" + ' + + - label: "CosyVoice3-TTS E2E Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - | + timeout 20m bash -c ' + pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model" + ' # TODO: Bagel test on ROCm is very unstable. @tjtanaa # Need to debug before reneable numerical changes across large PRs @@ -225,14 +248,3 @@ steps: # - export VLLM_WORKER_MULTIPROC_METHOD=spawn # - export VLLM_ROCM_USE_AITER_RMSNORM=0 # - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm" - -- label: "CosyVoice3-TTS E2E Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 20m bash -c ' - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model" - ' diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 191b7f214..f01a07260 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -206,7 +206,34 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "TTS · Ming-Omni-TTS Test" + timeout_in_minutes: 20 + commands: + - | + timeout 20m bash -c " + export VLLM_LOGGING_LEVEL=DEBUG + pytest -s -v tests/e2e/online_serving/test_ming_tts.py tests/e2e/offline_inference/test_ming_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model' + " + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "TTS · Voxtral-TTS Test" + source_file_dependencies: + - tests/e2e/offline_inference/test_voxtral_tts.py + - tests/e2e/online_serving/test_voxtral_tts.py + - vllm_omni/model_executor/models/voxtral_tts/ + - vllm_omni/model_executor/stage_input_processors/voxtral_tts.py + - vllm_omni/deploy/voxtral_tts.yaml timeout_in_minutes: 20 commands: - | @@ -249,6 +276,12 @@ steps: type: DirectoryOrCreate - label: "TTS · GLM-TTS Test" + source_file_dependencies: + - tests/e2e/online_serving/test_glm_tts.py + - tests/e2e/offline_inference/test_glm_tts.py + - vllm_omni/model_executor/models/glm_tts/ + - vllm_omni/model_executor/stage_input_processors/glm_tts.py + - vllm_omni/deploy/glm_tts.yaml timeout_in_minutes: 20 commands: - | @@ -291,6 +324,12 @@ steps: type: DirectoryOrCreate - label: "Omni · Qwen3-Omni Test" + source_file_dependencies: + - tests/e2e/offline_inference/test_qwen3_omni.py + - tests/e2e/online_serving/test_qwen3_omni.py + - vllm_omni/model_executor/models/qwen3_omni/ + - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py + - vllm_omni/deploy/qwen3_omni_moe.yaml timeout_in_minutes: 30 commands: - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' @@ -329,6 +368,11 @@ steps: type: DirectoryOrCreate - label: "Omni · Mimo Audio Test" + source_file_dependencies: + - tests/e2e/online_serving/test_mimo_audio.py + - vllm_omni/model_executor/models/mimo_audio/ + - vllm_omni/model_executor/stage_input_processors/mimo_audio.py + - vllm_omni/deploy/mimo_audio.yaml timeout_in_minutes: 30 commands: - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m 'advanced_model and cuda' --run-level 'advanced_model' @@ -366,37 +410,14 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: "Omni · VoxCPM2 Test" - timeout_in_minutes: 20 - commands: - - | - timeout 20m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'advanced_model and cuda' --run-level 'advanced_model' - " - agents: - queue: "gpu_1_queue" - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - - "HF_TOKEN" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - - label: "TTS · VoxCPM2 Online Test" + - label: "TTS · VoxCPM2 Test" timeout_in_minutes: 30 commands: - | timeout 30m bash -c " export VLLM_LOGGING_LEVEL=DEBUG export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model' + pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model' " agents: queue: "gpu_1_queue" @@ -438,6 +459,9 @@ steps: - "/fsx/hf_cache:/fsx/hf_cache" - label: "Diffusion · Qwen Image Edit Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image_edit.py + - vllm_omni/diffusion/models/qwen_image/ timeout_in_minutes: 20 commands: - pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model' @@ -476,6 +500,9 @@ steps: type: DirectoryOrCreate - label: "Diffusion · Longcat Image Edit Test" + source_file_dependencies: + - tests/e2e/online_serving/test_longcat_image_edit.py + - vllm_omni/diffusion/models/longcat_image/ timeout_in_minutes: 20 commands: - pytest -s -v tests/e2e/online_serving/test_longcat_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model' @@ -514,6 +541,9 @@ steps: type: DirectoryOrCreate - label: "Diffusion · Longcat Image Test" + source_file_dependencies: + - tests/e2e/online_serving/test_longcat_image.py + - vllm_omni/diffusion/models/longcat_image/ timeout_in_minutes: 20 commands: - pytest -s -v tests/e2e/online_serving/test_longcat_image.py -m 'advanced_model and cuda' --run-level 'advanced_model' @@ -552,6 +582,13 @@ steps: type: DirectoryOrCreate - label: "Diffusion · Bagel Test" + source_file_dependencies: + - tests/e2e/offline_inference/test_bagel.py + - tests/e2e/online_serving/test_bagel.py + - vllm_omni/model_executor/models/bagel/ + - vllm_omni/diffusion/models/bagel/ + - vllm_omni/model_executor/stage_input_processors/bagel.py + - vllm_omni/deploy/bagel.yaml timeout_in_minutes: 20 commands: - pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model' @@ -591,6 +628,9 @@ steps: - label: "Diffusion · Qwen Image Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image.py + - vllm_omni/diffusion/models/qwen_image/ commands: - | timeout 40m bash -c ' @@ -632,6 +672,9 @@ steps: type: DirectoryOrCreate - label: "Diffusion · Qwen Image Layered Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image_layered.py + - vllm_omni/diffusion/models/qwen_image/ commands: - | timeout 40m bash -c " @@ -673,6 +716,10 @@ steps: type: DirectoryOrCreate - label: "Diffusion · Wan22 Test" + source_file_dependencies: + - tests/e2e/offline_inference/test_wan22_t2v.py + - tests/e2e/online_serving/test_wan22_t2v.py + - vllm_omni/diffusion/models/wan2_2/ commands: - | timeout 40m bash -c " @@ -714,6 +761,9 @@ steps: type: DirectoryOrCreate - label: "Diffusion · HunyuanVideo 1.5 Test" + source_file_dependencies: + - tests/e2e/online_serving/test_hunyuan_video_15.py + - vllm_omni/diffusion/models/hunyuan_video/ timeout_in_minutes: 40 commands: - | @@ -789,6 +839,9 @@ steps: - "/fsx/hf_cache:/fsx/hf_cache" - label: "Diffusion · SD3 Test" + source_file_dependencies: + - tests/e2e/online_serving/test_sd3.py + - vllm_omni/diffusion/models/sd3/ timeout_in_minutes: 45 commands: - | diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index ef99257fb..dd6d5ec31 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -568,6 +568,55 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test" + key: nightly-tts-performance-higgs-audio-v3 + depends_on: upload-nightly-pipeline + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" + timeout_in_minutes: 180 + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - export VLLM_USE_DEEP_GEMM="0" + - export VLLM_MOE_USE_DEEP_GEMM="0" + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below. - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test" key: nightly-diffusion-x2iat-group @@ -591,6 +640,7 @@ steps: tests/e2e/online_serving/test_longcat_image_edit_expansion.py tests/e2e/online_serving/test_flux_2_dev_expansion.py tests/e2e/online_serving/test_bagel_expansion.py + tests/e2e/offline_inference/test_sensenova_u1_expansion.py -m "full_model and diffusion and H100 and not distributed_cuda" --run-level "full_model" agents: @@ -1462,6 +1512,7 @@ steps: - nightly-omni-performance-vllm-text - nightly-omni-performance-multi-replicas - nightly-tts-performance + - nightly-tts-performance-higgs-audio-v3 - nightly-diffusion-x2iat-performance-qwen-image - nightly-diffusion-x2iat-performance-qwen-image-edit - nightly-diffusion-x2iat-performance-qwen-image-edit-2511 @@ -1475,6 +1526,7 @@ steps: - export DEFAULT_INPUT_DIR=tests/dfx/perf/results - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3 - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 417db45be..699b5e94b 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -215,6 +215,43 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: "Entrypoints Test" + depends_on: upload-ready-pipeline + commands: + - | + timeout 30m bash -c " + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -sv tests/entrypoints/ -m 'core_model and cuda' --run-level 'core_model' + " + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - group: ":card_index_dividers: E2E Test" depends_on: upload-ready-pipeline steps: @@ -239,6 +276,11 @@ steps: - "/fsx/hf_cache:/fsx/hf_cache" - label: "Omni · Qwen3-Omni Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen3_omni.py + - vllm_omni/model_executor/models/qwen3_omni/ + - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py + - vllm_omni/deploy/qwen3_omni_moe.yaml commands: - | timeout 20m bash -c " @@ -273,73 +315,14 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: "Omni · Qwen3-Omni Realtime WebSocket" - commands: - - | - timeout 30m bash -c " - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -m 'core_model' --run-level 'core_model' - " - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: "Omni · VoxCPM2 Test" - timeout_in_minutes: 20 - commands: - - | - timeout 20m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'core_model' --run-level 'core_model' - " - agents: - queue: "gpu_1_queue" - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - - "HF_TOKEN" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - - label: "TTS · VoxCPM2 Online Test" + - label: "TTS · VoxCPM2 Test" timeout_in_minutes: 30 commands: - | timeout 30m bash -c " export VLLM_LOGGING_LEVEL=DEBUG export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model' + pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model' " agents: queue: "gpu_1_queue" @@ -399,7 +382,33 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "TTS · Higgs-Audio-V3 Test" + commands: + - | + timeout 20m bash -c " + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_USE_DEEP_GEMM=0 + export VLLM_MOE_USE_DEEP_GEMM=0 + pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model' + " + agents: + queue: "gpu_4_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Diffusion · Qwen Image Edit Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image_edit.py + - vllm_omni/diffusion/models/qwen_image/ commands: - | timeout 20m bash -c " @@ -435,6 +444,12 @@ steps: type: DirectoryOrCreate - label: "Diffusion · Bagel Test" + source_file_dependencies: + - tests/e2e/online_serving/test_bagel.py + - vllm_omni/model_executor/models/bagel/ + - vllm_omni/diffusion/models/bagel/ + - vllm_omni/model_executor/stage_input_processors/bagel.py + - vllm_omni/deploy/bagel.yaml commands: - | timeout 40m bash -c " @@ -476,6 +491,9 @@ steps: type: DirectoryOrCreate - label: "Diffusion · Wan22 Test" + source_file_dependencies: + - tests/e2e/online_serving/test_wan22_t2v.py + - vllm_omni/diffusion/models/wan2_2/ commands: - | timeout 40m bash -c " @@ -518,6 +536,9 @@ steps: - label: "Diffusion · Qwen Image Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image.py + - vllm_omni/diffusion/models/qwen_image/ commands: - | timeout 40m bash -c " @@ -559,6 +580,9 @@ steps: type: DirectoryOrCreate - label: "Diffusion · Qwen Image Layered Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image_layered.py + - vllm_omni/diffusion/models/qwen_image/ commands: - | timeout 40m bash -c " diff --git a/.buildkite/test-template-amd-omni.j2 b/.buildkite/test-template-amd-omni.j2 index db76f8216..aa5ca8bcd 100644 --- a/.buildkite/test-template-amd-omni.j2 +++ b/.buildkite/test-template-amd-omni.j2 @@ -6,6 +6,13 @@ {% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %} {% set default_working_dir = "/app/vllm-omni" %} +{% if env %} +env: +{% for name in env %} + {{ name }}: "{{ env[name] }}" +{% endfor %} +{% endif %} +steps: - group: "AMD Tests" depends_on: ~ steps: @@ -31,8 +38,15 @@ agents: queue: amd-cpu - {% for step in steps %} - {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %} +{% for step_group in steps %} +{% if step_group.group %} + - group: "{{ step_group.group }}" +{% if step_group.depends_on %} + depends_on: {{ step_group.depends_on }} +{% endif %} + steps: +{% for step in step_group.steps %} +{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %} - label: "{{ step.agent_pool }}: {{ step.label }}" depends_on: amd-build agents: @@ -46,8 +60,23 @@ command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh env: DOCKER_BUILDKIT: "1" +{% if step.env %} +{% for name in step.env %} + {{ name }}: "{{ step.env[name] }}" +{% endfor %} +{% endif %} TEST_COMMAND: |- (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} +{% if env %} +{% for name in env %} + export {{ name }}="{{ env[name] }}" +{% endfor %} +{% endif %} +{% if step.env %} +{% for name in step.env %} + export {{ name }}="{{ step.env[name] }}" +{% endfor %} +{% endif %} {% if "mi250" in step.agent_pool %} python3 -m pip uninstall -y amd-aiter {% endif %} @@ -61,5 +90,52 @@ {% else %} soft_fail: true {% endif%} +{% endif %} +{% endfor %} +{% elif step_group.mirror_hardwares and mirror_hw in step_group.mirror_hardwares %} +{% set step = step_group %} + - label: "{{ step.agent_pool }}: {{ step.label }}" + depends_on: amd-build + agents: + {% if step.agent_pool %} + queue: amd_{{ step.agent_pool }} + {% else %} + queue: amd_mi325_1 + {% endif %} +{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %} +{% set indented_cmd = cmd_body | replace("\n", "\n ") %} + command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh + env: + DOCKER_BUILDKIT: "1" +{% if step.env %} +{% for name in step.env %} + {{ name }}: "{{ step.env[name] }}" +{% endfor %} +{% endif %} + TEST_COMMAND: |- + (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} +{% if env %} +{% for name in env %} + export {{ name }}="{{ env[name] }}" +{% endfor %} +{% endif %} +{% if step.env %} +{% for name in step.env %} + export {{ name }}="{{ step.env[name] }}" +{% endfor %} +{% endif %} +{% if "mi250" in step.agent_pool %} + python3 -m pip uninstall -y amd-aiter +{% endif %} + {{ indented_cmd | safe }} + priority: 100 + {% if step.timeout_in_minutes %} + timeout_in_minutes: {{ step.timeout_in_minutes }} {% endif %} - {% endfor %} + {% if step.grade and step.grade == "Blocking" %} + soft_fail: false + {% else %} + soft_fail: true + {% endif%} +{% endif %} +{% endfor %} diff --git a/.claude/skills/precheck-pr/SKILL.md b/.claude/skills/precheck-pr/SKILL.md new file mode 100644 index 000000000..ab6034a48 --- /dev/null +++ b/.claude/skills/precheck-pr/SKILL.md @@ -0,0 +1,97 @@ +--- +name: precheck-pr +description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub. +--- + +# PR Pre-Check + +Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only. + +## Mode Selection + +| Mode | When | Time | +|------|------|------| +| **Quick** | About to push, final sanity check | ~3 min | +| **Full** | Ready for review, want maintainer-level scan | ~10 min | + +Default to quick if unsure. Run full before marking a PR "ready for review." + +## Workflow + +### Step 1: Detect Base Branch + +```bash +BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \ + || git merge-base HEAD main 2>/dev/null \ + || echo origin/main) +echo "diffing against ${BASE_SHA}" +git diff --name-only ${BASE_SHA}...HEAD +``` + +### Step 2: Validate PR Title + +Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes: + +| Prefix | Applies to | +|--------|-----------| +| `[Bugfix]` | Bug fixes | +| `[CI/Build]` | Build or CI improvements | +| `[Doc]` | Documentation changes | +| `[Model]` | New/improved models (include model name) | +| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) | +| `[Kernel]` | CUDA/kernel changes | +| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) | +| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) | +| `[Misc]` | Other changes (use sparingly) | + +✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title. +⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add ...`). + +### Step 3: Categorize the PR + +| Diff contains | PR type | +|---------------|---------| +| New files under `vllm_omni/model_executor/models//` | **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 + + Mode: quick | full + Type: + + 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. diff --git a/.claude/skills/precheck-pr/references/checklists.md b/.claude/skills/precheck-pr/references/checklists.md new file mode 100644 index 000000000..5b9472d96 --- /dev/null +++ b/.claude/skills/precheck-pr/references/checklists.md @@ -0,0 +1,135 @@ +# Pre-Check Checklists + +Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing. + +--- + +## All PRs: Title Format (check first) + +PR title must follow the project convention documented in `docs/contributing/README.md`: + +| Prefix | Applies to | +|--------|-----------| +| `[Bugfix]` | Bug fixes | +| `[CI/Build]` | Build or CI improvements | +| `[Doc]` | Documentation changes | +| `[Model]` | New/improved models (include model name) | +| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) | +| `[Kernel]` | CUDA/kernel changes | +| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) | +| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) | +| `[Misc]` | Other changes (use sparingly) | + +**Examples:** +``` +✓ [Model] Add MiniCPM-o 4.5 support +✓ [Bugfix] Fix Ovis image text encoder dtype +✓ [Core] Qwen3-Omni performance optimization +✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend +✗ Fix Ovis image text encoder dtype ← missing prefix +✗ [bugfix] Fix Ovis image text encoder ← wrong case +✗ [WIP] Add model ← WIP in title +``` + +- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add ...`). ⚠ 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 diff --git a/.claude/skills/quantization/SKILL.md b/.claude/skills/quantization/SKILL.md new file mode 100644 index 000000000..b58daaa28 --- /dev/null +++ b/.claude/skills/quantization/SKILL.md @@ -0,0 +1,103 @@ +--- +name: quantization +description: Work on vLLM-Omni quantization for diffusion, autoregressive, omni, or multi-stage models. Use when choosing or adding methods such as fp8, int8, gguf, mxfp8, mxfp4, mxfp4_dualscale, ModelOpt, AutoRound, INC, msModelSlim, awq, or gptq; debugging quantized loading; or validating memory, speed, and output quality. +--- + +# vLLM-Omni Quantization + +Use this skill for quantization work in `vllm-omni`. Start from the local +code and docs, then pick the smallest validated path for the model, method, +hardware, and quality target. + +## First Checks + +Before changing code or recommending a command, identify: + +- Model and task: AR text, omni/TTS, text-to-image, text-to-video, or multi-stage diffusion. +- Hardware backend: CUDA generation, Ascend NPU, Intel XPU, or another platform. +- Quantization mode: online from BF16/FP16 weights, pre-quantized checkpoint, runtime KV/FA quantization, or per-component routing. +- Scope: transformer/DiT, thinker language model, text encoder, VAE, audio/vision encoder, talker, or a specific stage. +- Validation target: memory reduction, latency/throughput, output quality, or all three. + +Do not assume a method is supported just because the CLI accepts the string. +Check `docs/user_guide/quantization/` and the closest implementation first. + +## Quick Decision + +| Task | Start With | +|------|------------| +| Choose a method or command | `references/methods.md` and `references/modality-compat.md` | +| Use `build_quant_config()` or per-component routing | `references/methods.md` | +| Work on diffusion quantization | `references/diffusion.md` | +| Add quantization to a new model | `references/adding-models.md` | +| Convert or load ModelOpt FP8 checkpoints | `references/modelopt-fp8.md` | +| Debug ModelOpt, GGUF, AutoRound, MXFP, or serialized Int8 loading | `references/diffusion.md` and method docs under `docs/user_guide/quantization/` | + +## Current Runtime Shape + +The unified entrypoint is: + +```python +from vllm_omni.quantization import build_quant_config +``` + +It supports method strings, flat method dictionaries, per-component +dictionaries, existing `QuantizationConfig` objects, and `None`. The factory +delegates generic methods to upstream `vllm` and keeps vLLM-Omni overrides for +diffusion or omni-specific routing: + +- `gguf` +- `int8` +- `mxfp8` +- `mxfp4` +- `mxfp4_dualscale` +- `inc`, `auto-round`, `auto_round` +- ModelOpt auto-detected configs: `modelopt`, `modelopt_fp4`, `modelopt_mixed` + +Use `ComponentQuantizationConfig` when only one stage or component should be +quantized. Pre-quantized ModelOpt-style checkpoints should not spill into +vision/audio encoders that have no corresponding scale tensors. + +## Ownership Boundary + +- Upstream `vllm` owns generic quantization configs, kernels, loader semantics, + hardware capability rules, and generic AR quantization methods. +- `vllm-omni` owns unified config routing, diffusion-specific wrappers, + component scoping, GGUF or ModelOpt checkpoint adapters, model-specific + prefix mapping, docs, examples, and validation. + +If a new method needs missing generic kernels or loader behavior, fix upstream +`vllm` first. In `vllm-omni`, add thin integration and model wiring. + +## Working Loop + +1. Read the local method doc under `docs/user_guide/quantization/`. +2. Inspect the active factory and config class in `vllm_omni/quantization/`. +3. Inspect the target model pipeline and transformer for stable prefixes and + whether every relevant vLLM linear layer receives `quant_config`. +4. For pre-quantized checkpoints, inspect `transformer/config.json` and the + checkpoint tensor names before running full generation. +5. Validate with the same prompt, seed, size, scheduler, steps, dtype, eager or + non-eager mode, and parallelism as the BF16 baseline. +6. Report quality, latency, and memory separately. Do not call a method + supported until quality is checked. + +## Common Mistakes + +| Symptom | Likely Cause | Fix | +|---------|--------------|-----| +| `--quantization` has no visible effect | Wrong scope or unsupported model path | Check component routing and method docs | +| Some layers stay BF16 unexpectedly | `quant_config` was not threaded into all vLLM linear layers | Audit transformer constructors and prefixes | +| Quality collapses but loading succeeds | Too many sensitive layers were quantized | Add model-specific `ignored_layers` and compare to BF16 | +| ModelOpt checkpoint loads but output is corrupted | Prefixes, packed-module mapping, scale routing, or BF16 fallback is wrong | Use `references/modelopt-fp8.md` | +| GGUF shape or tensor mismatch | Missing architecture-specific adapter | Add explicit adapter mapping; avoid generic fallback | +| Online and offline MXFP4 disagree | Wrong `mxfp4` vs `mxfp4_dualscale` mode | Check `docs/user_guide/quantization/mxfp4.md` | +| Quantized path is slower than BF16 | Kernel path, compile mode, dtype, or shape mismatch | Compare eager-to-eager and non-eager-to-non-eager | + +## References + +- General methods and config routing: [references/methods.md](references/methods.md) +- Modality and model scope: [references/modality-compat.md](references/modality-compat.md) +- Diffusion implementation patterns: [references/diffusion.md](references/diffusion.md) +- Adding model support: [references/adding-models.md](references/adding-models.md) +- ModelOpt FP8 conversion and loading: [references/modelopt-fp8.md](references/modelopt-fp8.md) diff --git a/.claude/skills/quantization/references/adding-models.md b/.claude/skills/quantization/references/adding-models.md new file mode 100644 index 000000000..ac9554de2 --- /dev/null +++ b/.claude/skills/quantization/references/adding-models.md @@ -0,0 +1,70 @@ +# Adding Quantization Support to a Model + +Use this checklist when extending quantization to a new architecture or stage. + +## Classify Ownership + +| Case | Owner | +|------|-------| +| New generic AR method, kernel, or loader semantics | upstream `vllm` first | +| Existing upstream method needs diffusion wiring | `vllm-omni` model and config integration | +| GGUF or ModelOpt tensor mapping for a diffusion architecture | `vllm-omni` adapter | +| NPU MXFP or msModelSlim diffusion path | `vllm-omni` platform/method wrapper and docs | +| Multi-stage omni checkpoint quantizes only the thinker LM | `vllm-omni` component scoping | + +Avoid building a private quantization stack in model code. + +## Model Wiring Checklist + +1. Find the closest supported model implementation. +2. Confirm the model uses vLLM linear layers that accept `quant_config`. +3. Add `quant_config` to the transformer constructor. +4. Pass stable `prefix` values through every linear layer. +5. Handle fused modules and packed mappings explicitly. +6. Ensure `load_weights()` and any `WeightsMapper` remap `ignored_layers` or + `modules_to_not_convert`. +7. Keep tokenizer, scheduler, VAE, text encoder, audio/vision encoders, and + talker BF16 unless method-specific docs validate them. +8. Add or update a method doc under `docs/user_guide/quantization/`. + +## Checkpoint Adapter Checklist + +For pre-quantized checkpoints: + +- Inspect `transformer/config.json` first. +- Check `quant_method`, `method`, `producer.name`, `quant_algo`, serialized + flags, `ignored_layers`, and any method-specific fields. +- Verify tensor name mapping before generation. +- Verify packed modules such as `to_qkv`, `add_kv_proj`, and `w13`. +- Confirm scale tensor names and shapes. +- Allow dequantization back to BF16 for intentionally skipped target layers. + +## Sensitive Layer Discovery + +Start conservative. Quantize the backbone first and keep output-side or routing +layers BF16 until quality is proven. + +Frequently sensitive areas: + +- attention output projections such as `to_out` +- FFN down projections such as `w2` +- final projections such as `proj_out`, `final_layer`, and output heads +- modulation, timestep, position, image, or text embedders +- refiner blocks +- MoE routers, experts selected by routing, and layernorm-adjacent modules +- VAE, text/vision towers, audio encoders, and codec/talker stages + +Do not reuse one model family's ignore list blindly. + +## Tests and Evidence + +Minimum evidence for a PR: + +- loader/config test or smoke command +- fixed-seed BF16 vs quantized output comparison +- quality metrics and saved outputs for precision-changing work +- memory and latency comparison under identical execution mode +- failure behavior for unsupported models or missing checkpoint fields + +When the PR adds a method to a support table, include the exact model, +hardware, command, quality threshold, and artifact path used for validation. diff --git a/.claude/skills/quantization/references/diffusion.md b/.claude/skills/quantization/references/diffusion.md new file mode 100644 index 000000000..1caeef0cf --- /dev/null +++ b/.claude/skills/quantization/references/diffusion.md @@ -0,0 +1,113 @@ +# Diffusion Quantization Reference + +Use this file for DiT quantization, config routing, checkpoint adapters, and +quality or loader debugging. + +## Supported Method Families + +The current local diffusion quantization families are: + +- CUDA/GPU online or checkpoint: `fp8`, `int8`, `gguf`, ModelOpt +- Ascend NPU: `int8`, `mxfp8`, `mxfp4`, `mxfp4_dualscale`, msModelSlim paths +- Intel XPU: AutoRound/INC and online MXFP8 paths where documented + +Always cross-check the method doc before claiming model support. The CLI +accepting a method string is not enough. + +## Implementation Starting Points + +| Area | Files | +|------|-------| +| Unified factory | `vllm_omni/quantization/factory.py` | +| Per-component routing | `vllm_omni/quantization/component_config.py` | +| GGUF diffusion method | `vllm_omni/quantization/gguf_config.py` | +| Int8 diffusion method | `vllm_omni/quantization/int8_config.py` | +| MXFP8 diffusion method | `vllm_omni/quantization/mxfp8_config.py` | +| MXFP4 and dual-scale method | `vllm_omni/quantization/mxfp4_config.py` | +| AutoRound/INC omni extension | `vllm_omni/quantization/inc_config.py` | +| ModelOpt checkpoint adapters | `vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py` | +| FP8 scaled-mm runtime patches | `vllm_omni/patch.py` | +| Similarity tool | `vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` | + +## Threading `quant_config` + +For a diffusion transformer: + +1. Accept `quant_config: QuantizationConfig | None` in the transformer. +2. Pass `quant_config` and stable `prefix` values to every vLLM linear layer. +3. Preserve packed-module mappings for fused QKV, FFN, or model-specific fused + projections. +4. Call `apply_vllm_mapper()` on configs that need HF-to-vLLM name remapping. +5. For component configs, keep prefixes specific enough to avoid quantizing VAE + or encoders by accident. + +Common skipped areas are `proj_out`, output projections, modulation layers, +embedders, refiner blocks, MoE routers, and layernorm-adjacent paths. + +## Online vs Offline + +Online quantization: + +- starts from BF16/FP16 checkpoint weights +- computes or stores quantized weights at load time +- typically uses `--quantization ` or `quantization=""` + +Offline or pre-quantized checkpoints: + +- include quantized weights and scales on disk +- usually carry `quantization_config` in `transformer/config.json` +- should auto-detect when the config is complete +- may require model-specific adapter logic for tensor names and packed modules + +Do not mix online flags with incompatible offline checkpoint configs. If the +active config and disk config differ, `resolve_quant_config_from_disk()` should +raise or rebuild rather than silently loading corrupted weights. + +## GGUF Rules + +- Require `quantization_config.gguf_model`. +- Use the normal base model for tokenizer, scheduler, text encoder, and VAE. +- Add explicit architecture adapters for tensor-name mapping. +- Use `named_parameters()` and `named_buffers()`, not `state_dict()`, to + discover loadable names. +- Guard fused QKV/KV rewrites so tensors are not rewritten twice. +- Flatten N-D activations around the GGUF matmul and restore shape. + +Unsupported architectures should fail clearly instead of falling back to a fake +generic mapper. + +## MXFP Rules + +- `mxfp8` is W8A8 microscaling, currently centered on NPU and documented XPU + behavior. +- `mxfp4` is online single-scale W4A4 on NPU. +- `mxfp4_dualscale` is online/offline dual-scale W4A4 with BF16 fallback; the + offline path is the production path when calibrated `mul_scale` exists. +- Offline dual-scale checkpoints must inject per-transformer `ignored_layers` + into config because cascade transformers can have different BF16 fallback + sets. + +Do not load an offline dual-scale checkpoint with the online single-scale path. + +## Common Failures + +| Symptom | Likely Cause | Fix | +|---------|--------------|-----| +| `Unknown quantization method` | Factory override or upstream registration is missing | Add the right override or upstream method | +| Missing scale tensors | Wrong scope, wrong adapter, or BF16 layer not allowed to dequantize | Check config scope and adapter scale routing | +| Output rank collapses from 3-D to 2-D | FP8 scaled-mm kernel ignored output shape | Check the patch in `vllm_omni/patch.py` | +| Quality degrades silently | Sensitive layers quantized | Add `ignored_layers`; compare to BF16 | +| Only one cascade transformer works | Disk config differs per transformer but active config is reused | Rebuild from each transformer's config | +| Quantized path slower than BF16 | Kernel, dtype, compile, or benchmark mode mismatch | Compare eager-to-eager and non-eager-to-non-eager | + +## Verification + +For a new or changed diffusion quantization path, run: + +1. Config and loader smoke test. +2. One fixed-seed generation against BF16. +3. Numeric comparison with the similarity tool or equivalent metrics. +4. Serving test if the path is user-facing online. +5. Memory and latency check under the same execution mode as the baseline. + +Keep output artifacts when precision changes are part of the PR. diff --git a/.claude/skills/quantization/references/methods.md b/.claude/skills/quantization/references/methods.md new file mode 100644 index 000000000..4a1e0a9b8 --- /dev/null +++ b/.claude/skills/quantization/references/methods.md @@ -0,0 +1,82 @@ +# General Quantization Methods + +Use this file for method selection and unified `vllm-omni` config syntax. For +diffusion-specific wiring, see `diffusion.md`. For ModelOpt checkpoint export +and loading details, see `modelopt-fp8.md`. + +## Unified Entry Point + +`vllm_omni.quantization.build_quant_config()` accepts: + +- `None` or `"none"` to disable quantization +- a method string such as `"fp8"`, `"int8"`, `"mxfp8"`, `"awq"`, or `"auto-round"` +- a flat dict with `"method"` or `"quant_method"` +- a per-component dict such as `{"transformer": {"method": "fp8"}, "vae": None}` +- an already-built `QuantizationConfig` + +Per-component dicts are routed by `ComponentQuantizationConfig` using +longest-prefix matching. Prefix stability matters: if a model remaps HF names +to vLLM names, apply the mapper to `ignored_layers` or component prefixes. + +## Local Overrides + +The local factory delegates most methods to upstream `vllm`, but has vLLM-Omni +overrides for methods that need diffusion, NPU, XPU, or component behavior: + +| Method | Local Config | Main Use | +|--------|--------------|----------| +| `gguf` | `DiffusionGGUFConfig` | Pre-quantized diffusion transformer weights | +| `int8` | `DiffusionInt8Config` | Online or serialized W8A8 diffusion transformers | +| `mxfp8` | `DiffusionMXFP8Config` | W8A8 MXFP8 diffusion on NPU and online XPU path | +| `mxfp4` | `DiffusionMXFP4Config` | Online W4A4 MXFP4 diffusion on NPU | +| `mxfp4_dualscale` | `DiffusionMXFP4DualScaleMixedConfig` | Online/offline W4A4 dual-scale plus BF16 fallback on NPU | +| `inc`, `auto-round`, `auto_round` | `OmniINCConfig` | AutoRound/INC checkpoint loading, including MXFP8 checkpoint metadata | +| ModelOpt config dicts | upstream ModelOpt configs through auto-detect | FP8, NVFP4, or mixed ModelOpt checkpoints | + +Check `SUPPORTED_QUANTIZATION_METHODS` in `vllm_omni/quantization/factory.py` +before adding a new method name. + +## AR and Generic Quantization + +For AR-backed models, upstream `vllm` methods usually apply: + +- `awq` and `gptq`: 4-bit or 8-bit weight quantization with calibration, + commonly used for text and AR stages. +- `fp8`: FP8 weight path or KV-cache FP8 depending on flags and hardware. +- `inc` or `auto-round`: pre-quantized checkpoints with Intel Neural + Compressor / AutoRound metadata. + +Use these for the AR language model stage only when the checkpoint and model +guide support it. Omni and TTS models often keep audio, vision, talker, and +codec stages in BF16. + +## Diffusion and Offline Methods + +For diffusion transformers, prefer the method-specific docs: + +- `docs/user_guide/quantization/fp8.md` +- `docs/user_guide/quantization/int8.md` +- `docs/user_guide/quantization/gguf.md` +- `docs/user_guide/quantization/modelopt.md` +- `docs/user_guide/quantization/mxfp8.md` +- `docs/user_guide/quantization/mxfp4.md` +- `docs/user_guide/quantization/autoround.md` +- `docs/user_guide/quantization/msmodelslim.md` + +Pre-quantized checkpoints should normally auto-detect from config. Runtime +flags such as `--force-cutlass-fp8`, `--linear-backend cutlass`, and +`--moe-backend cutlass` select kernels; they do not quantize a BF16 checkpoint. + +## Validation + +Use fixed-seed A/B comparisons and record: + +- model path and quantized checkpoint path +- method config and `ignored_layers` +- prompt, size, frame count, steps, scheduler, dtype, and seed +- eager/non-eager mode, compile flags, and parallelism +- output metrics such as PSNR, MAE, cosine similarity, SSIM, LPIPS, plus visual artifacts +- latency and peak memory, reported separately from quality + +The helper `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity` +is the default lightweight image/video comparison tool for diffusion paths. diff --git a/.claude/skills/quantization/references/modality-compat.md b/.claude/skills/quantization/references/modality-compat.md new file mode 100644 index 000000000..188850e6b --- /dev/null +++ b/.claude/skills/quantization/references/modality-compat.md @@ -0,0 +1,47 @@ +# Quantization Compatibility by Modality + +Use this as a routing guide. The source of truth is the local method docs under +`docs/user_guide/quantization/` and the current model implementation. + +## Summary + +- AR and omni language-model quantization mostly follows upstream `vllm`, with + vLLM-Omni scoping for multi-stage models. +- Diffusion quantization is method-specific and model-specific. It usually + targets the diffusion transformer only. +- Pre-quantized checkpoints should quantize only the component that has + quantized weights and scale tensors. + +## Model Type Matrix + +| Model Type | Typical Quantization Scope | Notes | +|------------|----------------------------|-------| +| AR text models | Language model | Upstream `vllm` methods usually apply | +| Omni/TTS models | Thinker or AR language-model stage | Audio, vision, talker, and codec stages usually stay BF16 | +| DiT image models | Diffusion transformer | Check FP8, Int8, GGUF, ModelOpt, AutoRound docs | +| DiT video models | Diffusion transformer | HunyuanVideo FP8 and Wan2.2 MXFP paths are documented; validate other video paths | +| Multi-stage diffusion | Selected stage only | Use per-component routing and per-stage validation | +| Audio diffusion | Usually unsupported | Prefer dtype, offload, cache, or parallelism unless docs say otherwise | + +## Method Hints + +| Goal | Candidate | +|------|-----------| +| AR memory reduction | AWQ, GPTQ, AutoRound, ModelOpt where checkpoint supports it | +| CUDA diffusion W8A8 | FP8 or Int8 | +| Pre-quantized CUDA diffusion checkpoint | ModelOpt or GGUF | +| Ascend NPU diffusion | Int8, MXFP8, MXFP4, MXFP4 DualScale, msModelSlim | +| Intel XPU checkpoint | AutoRound/INC paths | +| Conservative image quality | selective FP8/Int8 with `ignored_layers` | +| Unsupported diffusion model | cache acceleration, CPU/layerwise offload, dtype, TP/SP/CFG parallelism | + +## Validation Rule + +Compatibility means more than successful loading. A method is compatible only +after it has: + +- correct component scope +- complete checkpoint/config metadata if offline +- fixed-seed quality comparison against BF16 +- acceptable latency and memory behavior +- docs or examples that state model, method, hardware, and caveats diff --git a/.claude/skills/quantization/references/modelopt-fp8.md b/.claude/skills/quantization/references/modelopt-fp8.md new file mode 100644 index 000000000..3c500e393 --- /dev/null +++ b/.claude/skills/quantization/references/modelopt-fp8.md @@ -0,0 +1,122 @@ +# ModelOpt FP8 Conversion and Loading + +Use this when exporting or loading ModelOpt FP8 checkpoints for diffusion +transformers. This reference captures the practical conversion lessons from the +2026-04-21 ModelOpt FP8 paste and aligns them with the current +`vllm-omni` ModelOpt adapter. + +## Core Rule + +Treat ModelOpt FP8 as selective checkpoint quantization, not one-click +whole-model FP8. Quantize the main transformer or DiT backbone, keep sensitive +or peripheral modules BF16, serialize a complete ModelOpt checkpoint, and let +vLLM-Omni auto-detect it from config. + +Do not pass `--quantization fp8` to force a ModelOpt checkpoint path. The +checkpoint's `quantization_config` should select `modelopt`. + +## Conversion Strategy + +Prefer: + +- backbone attention Q/K/V and FFN expansion projections +- conservative per-model ignore lists +- MHA quantizers disabled until basic linear quantization is stable +- full module-name prefixes in ignore/routing logic +- BF16 fallback for output-side and routing-sensitive layers + +Avoid quantizing by default: + +- VAE +- text, vision, or audio encoders +- embedders and modulation layers +- final heads and output projections +- refiner modules +- MoE routers and layernorm-sensitive blocks + +## Model-Specific Starting Points + +| Model Family | Quantize First | Keep BF16 First | +|--------------|----------------|-----------------| +| Z-Image | attention q/k/v, FFN w1/w3 | `to_out`, `w2`, `noise_refiner`, `context_refiner`, embedders, modulation, final/output | +| Qwen-Image | transformer backbone linear layers | `proj_out`, `img_in`, `txt_in`, norms, modulation, position/time/embed layers | +| FLUX2-dev / FLUX2-klein | backbone linear layers | `proj_out`, context/x embedders, `norm_out`, stream modulation, time embeddings | +| HunyuanImage-3.0 | `model.model` backbone | VAE, vision tower, token/time/image embedders, final image head, LM head, routers, layernorm-sensitive blocks | + +These are starting points, not support claims. Validate with BF16 outputs. + +## Required Checkpoint Metadata + +`transformer/config.json` must include a complete `quantization_config`. +At minimum it should identify ModelOpt and the algorithm: + +```json +{ + "quantization_config": { + "producer": {"name": "modelopt"}, + "quant_method": "modelopt", + "quant_algo": "FP8" + } +} +``` + +Current vLLM-Omni auto-detects: + +- `FP8` and `FP8_PER_CHANNEL_PER_TOKEN` -> `modelopt` +- `NVFP4` -> `modelopt_fp4` +- `MIXED_PRECISION` -> `modelopt_mixed` + +Packed modules and scales must map correctly: + +- `to_qkv` from `to_q`, `to_k`, `to_v` +- `add_kv_proj` from `add_q_proj`, `add_k_proj`, `add_v_proj` +- `w13` from `w1`, `w3` +- `.input_scale`, `.weight_scale`, `.weight_scale_2`, `.weight_scale_inv` + +For selective FP8, target layers that remain BF16 must be allowed to +dequantize from FP8 weights using the matching `weight_scale`, or the adapter +must skip the quantized scale tensors for that BF16 target. + +## Runtime Signals + +Healthy loading should show ModelOpt auto-detection and a ModelOpt linear +kernel path. Useful signals include: + +- auto-detected quantization from config +- selected ModelOpt FP8 linear method +- CUTLASS or FlashInfer FP8 scaled-mm backend when enabled and supported +- non-eager path compiling the transformer when expected + +If online serving fails while offline inference works, inspect warmup and +serving initialization before changing checkpoint conversion. + +## Validation Order + +1. Check config completeness. +2. Run an export checker such as `examples/quantization/check_modelopt_fp8_export.py`. +3. Confirm actual FP8 tensors and expected disk-size reduction. +4. Run one offline generation with fixed prompt, seed, size, and steps. +5. Run serving with `vllm serve --omni` and verify `/v1/images/generations`. +6. Compare against BF16 using visual review plus metrics such as LPIPS, PSNR, + MAE, cosine similarity, and SSIM. +7. Compare performance only under matched modes: + - BF16 eager vs FP8 eager + - BF16 non-eager vs FP8 non-eager + +FP8 is not automatically faster. Check whether execution reaches FP8 kernels, +whether compile mode is comparable, and whether the workload shape fits the +kernel path. + +## Triage + +| Symptom | First Check | +|---------|-------------| +| Startup failure | `config.json`, `model_index.json`, `quantization_config`, missing scales | +| Loads but output is corrupted | Over-quantized sensitive layers or wrong full-name prefixes | +| Offline works but online fails | Dummy warmup, serving init, compile path | +| FP8 is not faster | Kernel selection, dtype, compile mode, workload shape | +| BF16 fallback fails | Missing `weight_scale` for dequantization or wrong packed-module mapping | + +The success criterion is stable, correctly scoped checkpoint loading with +quality preserved enough for the target use case, not maximizing the number of +FP8 layers. diff --git a/.claude/skills/readme.md b/.claude/skills/readme.md index 44a4cd11d..f42fcb6ab 100644 --- a/.claude/skills/readme.md +++ b/.claude/skills/readme.md @@ -21,11 +21,18 @@ include: - `diffusion-perf-opt`: guides diffusion model performance optimization, including profiling traces, parallel strategies, stage timing analysis, and benchmark-driven tuning +- `quantization`: guides quantization method selection, model integration, + checkpoint loading, and quality/performance validation for vLLM-Omni - `add-omni-model`: covers addition of new omni-modality model support - `add-tts-model`: covers integration of new TTS models and related serving workflows - `generate-release-note`: helps prepare release notes for repository changes +- `precheck-pr`: self-check a branch before creating a PR — validates PR title + format, catches dead code, verifies accuracy/perf claims, and confirms merge + readiness - `review-pr`: provides a structured workflow for reviewing pull requests +- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the + latest vllm-ascend NPUModelRunner ## Maintenance Guidelines diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 1b4762794..adac6cf8a 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -1,21 +1,15 @@ -PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED. +PLEASE FILL IN THE PR DESCRIPTION HERE. ## Purpose ## Test Plan +**vLLM Version:** + +**vLLM-Omni Commit:** + ## Test Result ---- -
- Essential Elements of an Effective PR Description Checklist - -- [ ] 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. -
- -**BEFORE SUBMITTING, PLEASE READ ** (anything written below this line will be removed by GitHub Actions) +**BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions. +(anything written below this line will be removed by GitHub Actions) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index aac749775..4ace89fd6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,3 +1,5 @@ # Contributing to vLLM-Omni You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/) + +Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions. diff --git a/README.md b/README.md index 7dfc9cb01..cf1658b45 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone --- *Latest News* 🔥 +- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni). - [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends. - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments. - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup! @@ -56,8 +57,9 @@ vLLM-Omni is flexible and easy to use with: vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including: -- Omni-modality models (e.g. Qwen-Omni) -- Multi-modality generation models (e.g. Qwen-Image) +- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL) +- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3) +- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX) ## Getting Started diff --git a/benchmarks/accuracy/common.py b/benchmarks/accuracy/common.py index bce0b2359..bd28303b4 100644 --- a/benchmarks/accuracy/common.py +++ b/benchmarks/accuracy/common.py @@ -151,6 +151,9 @@ 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] @@ -170,6 +173,12 @@ 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 = [ ( diff --git a/benchmarks/accuracy/image_to_image/gedit_bench.py b/benchmarks/accuracy/image_to_image/gedit_bench.py index adb072c0a..2991bd4b5 100644 --- a/benchmarks/accuracy/image_to_image/gedit_bench.py +++ b/benchmarks/accuracy/image_to_image/gedit_bench.py @@ -64,6 +64,47 @@ 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) @@ -73,10 +114,7 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]: parsed = {"score": json.loads(stripped), "reasoning": ""} else: raise - score = parsed.get("score", []) - if not isinstance(score, list): - score = [score] - parsed["score"] = [int(value) for value in score] + parsed["score"] = _normalize_score_list(parsed.get("score", [])) return parsed @@ -391,6 +429,9 @@ 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 @@ -400,6 +441,9 @@ 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( @@ -478,6 +522,9 @@ 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 { @@ -633,6 +680,25 @@ 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) @@ -675,6 +741,9 @@ 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, diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 0a26bc0cb..9acf4e47e 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -33,6 +33,15 @@ models: voice_clone: task_type: Base + bosonai/higgs-audio-v3-tts-4b: + supported_tasks: [default_voice, voice_clone] + backend: openai-audio-speech + endpoint: /v1/audio/speech + trust_remote_code: true + task_extra_body: + default_voice: {} + voice_clone: {} + openbmb/VoxCPM2: supported_tasks: [voice_clone, default_voice] trust_remote_code: true diff --git a/docs/.nav.yml b/docs/.nav.yml index 562a49e84..258224e2a 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -87,6 +87,8 @@ nav: - AutoRound: user_guide/quantization/autoround.md - msModelSlim: user_guide/quantization/msmodelslim.md - ComfyUI: features/comfyui.md + - Fault Tolerance: + - Failure Modes: user_guide/fault_injection_reliability_matrix.md - Developer Guide: - General: - contributing/README.md diff --git a/docs/api/README.md b/docs/api/README.md index dce02e13d..d4c580629 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -92,7 +92,6 @@ Core scheduling and caching components. Configuration classes. - [vllm_omni.config.model.OmniModelConfig][] -- [vllm_omni.config.stage_config.ModelPipeline][] - [vllm_omni.config.stage_config.StageConfig][] - [vllm_omni.config.stage_config.StageConfigFactory][] - [vllm_omni.config.stage_config.StageType][] diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg index c4e836f68..2596240c1 100644 Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 8a5bcfff0..e1720b96e 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -116,6 +116,16 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat !!! note If the PR spans more than one category, please include all relevant prefixes. +### Pre-Check Before Submitting + +Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-review against project conventions: + +The skill offers two modes: +- **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status +- **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene + +The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) for the full checklist. + ### Local Test Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases. diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index ab0ae7f92..2efa39633 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -268,6 +268,101 @@ Before entering specific testing levels, the project establishes two common spec 1. ***PR Checklist ([Tests Style](../ci/tests_style.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance. 2. ***CI Failure Explanation ([CI Failures](../ci/failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency. +<<<<<<< ready +## Notes +### Diff-aware Buildkite uploads (`source_file_dependencies`) + +L2 (`.buildkite/test-ready.yml`) and L3 (`.buildkite/test-merge.yml`) pipelines can **skip unrelated GPU jobs at upload time** based on the PR diff. This is implemented by `.buildkite/scripts/upload_pipeline.py`, which filters steps before calling `buildkite-agent pipeline upload`. + +#### What `source_file_dependencies` is + +- A **uploader-only** YAML key on a Buildkite step or group. **Buildkite does not understand it**; `upload_pipeline.py` always **removes** it from the YAML that is uploaded. +- A list of path **prefixes** (directories or individual files). If **any** changed file in the diff equals a prefix or starts with `prefix/`, the step is kept; otherwise the step (or entire group) is **omitted** from the uploaded pipeline. + +#### When filtering runs + +| Build context | Changed files used for matching | +| --- | --- | +| Pull request | `git diff --name-only origin/...` | +| `main` branch push | `git diff --name-only ^..` | +| Other (e.g. local dry-run, non-PR branch) | Diff cannot be resolved → **no filtering**; all steps are uploaded and `source_file_dependencies` is still stripped | + +Docs-only PRs are handled earlier in bootstrap (`pipeline.yml`) via skip-ci logic; `source_file_dependencies` applies only to the **uploaded** L2/L3 test pipelines. + +#### Where it is configured + +| Level | Pipeline file | Upload entry | +| --- | --- | --- | +| L2 | `.buildkite/test-ready.yml` | `upload_pipeline.py --upload .buildkite/test-ready.yml` (from `pipeline.yml`) | +| L3 | `.buildkite/test-merge.yml` | `upload_pipeline.py --upload .buildkite/test-merge.yml` (from `pipeline.yml`) | + +Steps **without** `source_file_dependencies` are always uploaded (subject to the usual label conditions: `ready` for L2, `merge-test` for L3). + +#### Current skip policy (L2 / L3) + +To balance CI cost and coverage: + +- **Always run** (no `source_file_dependencies`): baseline groups such as Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3), and **all E2E jobs on L4 docker queues** (`gpu_1_queue`, `gpu_4_queue`). +- **Diff-gated** (`source_file_dependencies` set): **H100 E2E jobs only** (`queue: mithril-h100-pool` under the E2E Test group)—typically the model’s pytest file(s) plus related `vllm_omni/` model, deploy, and stage paths. + +Adding a new expensive H100 E2E step: list the test file(s) and the smallest set of source prefixes that should trigger it. Prefer **per-step** deps on leaf jobs rather than a broad group-level list unless every child shares the same paths. + +#### YAML example + +```yaml + - label: "Diffusion · Qwen Image Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image.py + - vllm_omni/diffusion/models/qwen_image/ + commands: + - pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' ... + agents: + queue: "mithril-h100-pool" +``` + +A **group** may also define `source_file_dependencies`; nested steps inherit filtering as a unit—the whole group is dropped if no prefix matches. + +#### Local dry-run + +```bash +# Render filtered YAML to stdout (no upload) +python3 .buildkite/scripts/upload_pipeline.py .buildkite/test-ready.yml + +# Confirm uploader-only keys are stripped +python3 .buildkite/scripts/upload_pipeline.py .buildkite/test-merge.yml | grep source_file_dependencies +# (no output expected) +``` + +On a PR build, Buildkite logs from `upload_pipeline.py` include lines such as `skip '…' (no changes under …)` for omitted steps. + +#### Related + +- Implementation: `.buildkite/scripts/upload_pipeline.py` +- L2/L3 diff skip does **not** replace label-based triggers (`ready`, `merge-test`); it only reduces which steps appear **after** the pipeline is already scheduled. +======= +### Test helper environment variables + +Some shared helpers under `tests/helpers/` honor optional environment variables for local debugging. These are **not** set in CI by default. + +| Variable | Accepted values | Description | +| -------- | --------------- | ----------- | +| `VLLM_OMNI_KEEP_REQUEST_MEDIA` | `1`, `true`, `yes` (case-insensitive) | When enabled, temporary WAV files created by `tests.helpers.media.convert_audio_bytes_to_text` are **not** deleted when the pytest process exits. By default, each call writes a unique file under the system temp directory via `tempfile.mkstemp` and registers `atexit` cleanup. Use this when debugging audio output validation (Whisper transcription, keyword checks, text–audio similarity). The saved path is logged as `audio data is saved: `. | + +Example (Linux / macOS): + +```bash +export VLLM_OMNI_KEEP_REQUEST_MEDIA=1 +pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -k test_mix_to_text_audio +``` + +Example (Windows PowerShell): + +```powershell +$env:VLLM_OMNI_KEEP_REQUEST_MEDIA = "1" +pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -k test_mix_to_text_audio +``` +>>>>>>> main + ## Chapter 1: L1 & L2 Level Testing - Unit Testing and Basic End-to-End Verification ### 1.1 Testing Purpose @@ -302,6 +397,7 @@ A clear directory structure is key to managing test cases efficiently. ### 1.4 Execution Method and Example - ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered. +- ***Diff-aware step skipping***: On L2, expensive **H100 E2E** jobs may be omitted at pipeline upload when the PR diff does not touch their [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) paths; L4 docker jobs and non-E2E steps still run. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies). - ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment. - ***Script Example***: @@ -404,7 +500,8 @@ L3 level testing executes after code is merged into the main branch. Its core pu ### 2.4 Execution Method and Example -- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch. +- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch (typically via `merge-test` label on the PR before merge). +- ***Diff-aware step skipping***: Same [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) mechanism as L2—only **H100 E2E** steps are diff-gated; L4 docker E2E and other L3 groups always upload. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies). - ***Execution Environment***: ***GPU*** servers. - ***Script Example***: @@ -520,6 +617,8 @@ L3 level testing executes after code is merged into the main branch. Its core pu **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation. + **Audio output debugging**: Deep validation may transcribe returned audio via `convert_audio_bytes_to_text` (Whisper). If an audio keyword or text–audio similarity assertion fails, set `VLLM_OMNI_KEEP_REQUEST_MEDIA=1` before running pytest to keep the intermediate WAV files for inspection (see [Test helper environment variables](#test-helper-environment-variables)). + - ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model` - ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model` diff --git a/docs/contributing/ci/failures.md b/docs/contributing/ci/failures.md index 68e468be1..80699e6a5 100644 --- a/docs/contributing/ci/failures.md +++ b/docs/contributing/ci/failures.md @@ -69,6 +69,8 @@ If the failure is already tracked, leave a comment on your PR noting that the fa If the failure appears to be new, investigate whether your changes caused it. +For resilience behavior under injected process signals/OOM pressure (validation scenarios, not root-cause equivalence), see [Failure Mode Injection Scenarios and Expected Behavior Matrix](../../user_guide/fault_injection_reliability_matrix.md). + ### Reading Buildkite Logs 1. Click **Details** next to the Buildkite check on your PR. diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index 755d6aeea..15d9e09df 100644 --- a/docs/design/feature/disaggregated_inference.md +++ b/docs/design/feature/disaggregated_inference.md @@ -8,6 +8,7 @@ Backend-specific setup lives in separate docs: - [SharedMemoryConnector](omni_connectors/shared_memory_connector.md) - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md) - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md) +- [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md) - [YuanrongConnector](omni_connectors/yuanrong_connector.md) - [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md) @@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode. | Single node | SharedMemoryConnector | Auto-configured if no connector is specified. | | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. | | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. | +| Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. | | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. | | Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. | diff --git a/docs/design/feature/omni_connectors/mori_transfer_engine_connector.md b/docs/design/feature/omni_connectors/mori_transfer_engine_connector.md new file mode 100644 index 000000000..73212c6de --- /dev/null +++ b/docs/design/feature/omni_connectors/mori_transfer_engine_connector.md @@ -0,0 +1,76 @@ +# MoriTransferEngineConnector + +## When to Use + +Currently supports intra-node deployment with Mori. + +As noted in #1742, inter-node support will be added back in a future +refactor. + +## Mechanism + +Uses Mori's `IOEngine` / `MemoryDesc` API for zero-copy RDMA transfers. + +- Data Plane: RDMA (InfiniBand/RoCE) with managed memory pool. +- Control Plane: ZMQ for pull-request handshake and async completion. + +## Installation + +See the [Mori repository](https://github.com/ROCm/mori) for installation instructions. + +## Configuration + +Mori is configured through the new deploy-config schema (see +[`docs/configuration/stage_configs.md`](../../../configuration/stage_configs.md)). +Define the connector at the top level of the deploy YAML and reference it +by name from each stage's `input_connectors` / `output_connectors`: + +```yaml +connectors: + mori_connector: + name: MoriTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + device_name: "" + memory_pool_size: 536870912 + memory_pool_device: "cuda" + +stages: + - stage_id: 0 + output_connectors: + to_stage_1: mori_connector + + - stage_id: 1 + input_connectors: + from_stage_0: mori_connector +``` + +A ready-to-run intra-node example for Qwen3-Omni-MoE on AMD MI300X lives +at `vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml` and can be loaded +with: + +```bash +vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \ + --deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml +``` + +The yaml wires `MoriTransferEngineConnector` (with `backend_type: xgmi`) to +the chunk_transfer_adapter path (`async_chunk: true`) so stage-to-stage +hidden-state and codec-frame streams ship GPU-to-GPU over AMD Infinity +Fabric instead of SHM. Qwen2.5-Omni + Mori is not yet functional on +the chunk path: it needs `thinker2talker_async_chunk` / +`talker2code2wav_async_chunk` input processors that do not exist yet +(the orchestrator-level path the upstream PR originally targeted was +lost during the entrypoints → engine refactor and removed in #1742). + +Parameters: + +- host: local RDMA IP (`"auto"` for auto-detect). +- zmq_port: ZMQ base port for control-plane communication. +- device_name: RDMA device (e.g., `"mlx5_0"`), empty for auto-detect. +- memory_pool_size: RDMA memory pool size in bytes. +- memory_pool_device: `"cpu"` (pinned) or `"cuda"` (GPUDirect / XGMI RDMA). + +For more details, refer to the +[Mori repository](https://github.com/ROCm/mori). diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index 009df77ff..50890c597 100644 --- a/docs/design/module/dit_module.md +++ b/docs/design/module/dit_module.md @@ -305,8 +305,8 @@ class GPUWorker: ) # Load model - model_loader = DiffusersPipelineLoader(load_config) - self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}") + model_loader = DiffusersPipelineLoader(load_config, self.od_config) + self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}") # Setup cache backend from vllm_omni.diffusion.cache.selector import get_cache_backend diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 7f7607c82..a8f293afd 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -45,7 +45,7 @@ If you do not need to modify source code of vLLM, you can directly install the s uv pip install vllm==0.22.0 --torch-backend=auto ``` -The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. +The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -70,8 +70,8 @@ git checkout v0.22.0 ``` Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file. ```bash -#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl +#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 4e3021f92..005c3f9e9 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -22,7 +22,7 @@ source .venv/bin/activate uv pip install vllm==0.22.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721 +uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index a0588a5c3..54b03afc8 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -33,7 +33,7 @@ th { | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | -| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | | +| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | @@ -50,6 +50,8 @@ th { | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | +| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | +| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | @@ -61,7 +63,12 @@ th { | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | +| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | +| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | +| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | +| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | @@ -69,11 +76,13 @@ th { | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | | | `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | +| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | +| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | |`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 61de36938..935deba4f 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -114,9 +114,10 @@ Content-Type: application/json | `instructions` | string | "" | Voice style/emotion instructions | | `max_new_tokens` | integer | 2048 | Maximum tokens to generate | | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. | +| `non_streaming_mode` | bool | null | Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. | | `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) | -**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian +**Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian. #### Voice Clone Parameters (Base task) @@ -332,7 +333,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "input": "Hello, this is a cloned voice", "task_type": "Base", "ref_audio": "https://example.com/reference.wav", - "ref_text": "Original transcript of the reference audio" + "ref_text": "Original transcript of the reference audio", + "non_streaming_mode": true }' --output cloned.wav ``` @@ -454,6 +456,7 @@ Content-Type: application/json | `ref_audio` | string | null | Default reference audio (Base task) | | `ref_text` | string | null | Default reference transcript (Base task) | | `max_new_tokens` | integer | null | Default max tokens | +| `non_streaming_mode` | bool | null | Default Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. | Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set: @@ -469,6 +472,7 @@ Each item in the `items` array requires only `input` (the text). All other field | `ref_audio` | string | Override reference audio | | `ref_text` | string | Override reference transcript | | `max_new_tokens` | integer | Override max tokens | +| `non_streaming_mode` | bool | Override Qwen3-TTS prompt construction mode. Does not affect HTTP response streaming or async-chunk pipelining. When null, inherits the batch-level value (then the model default). | ### Response Format diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index a8caee013..4d156bdf0 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -67,8 +67,10 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4 | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `input_reference` | file | null | Uploaded reference image for image-to-video requests | -| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` | +| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests | +| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` | +| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` | +| `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs | | `width` | integer | model default | Output video width | | `height` | integer | model default | Output video height | | `num_frames` | integer | model default | Number of generated frames | @@ -80,6 +82,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4 | `flow_shift` | number | null | Scheduler flow-shift value | | `true_cfg_scale` | number | null | True CFG scale when supported by the model | | `seed` | integer | null | Random seed for reproducibility | +| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation | +| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration | | `negative_prompt` | string | null | Text describing what to avoid in the generated video | | `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation | | `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on | @@ -123,6 +127,60 @@ curl -s http://localhost:8091/v1/videos \ -F "fps=16" ``` +### Video-to-Video + +For models that support video conditioning, upload the reference video with +`input_reference`: + +```bash +curl -s http://localhost:8091/v1/videos \ + -F "prompt=continue this motion with consistent subjects and lighting" \ + -F "input_reference=@input.mp4;type=video/mp4" \ + -F "width=1280" \ + -F "height=720" \ + -F "num_frames=80" \ + -F "fps=16" +``` + +You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload +through `video_reference`. Do not send `video_reference` together with +`input_reference` or `image_reference`. + +```bash +curl -s http://localhost:8091/v1/videos \ + -F "prompt=continue this motion with consistent subjects and lighting" \ + -F 'video_reference={"video_url":"https://example.com/input.mp4"}' \ + -F "width=1280" \ + -F "height=720" \ + -F "num_frames=80" \ + -F "fps=16" +``` + +JSON references currently support `image_url`/`video_url`; `file_id` references +are not implemented yet. Models may expose additional V2V controls through +`extra_params`. For example, Cosmos3 supports +`condition_frame_indexes_vision` and `condition_video_keep` to select which +decoded reference frames are used as clean conditioning. + +### Speech-to-Video + +For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both +an image reference and an audio reference. The `audio_reference` field accepts a +JSON string with `audio_url` pointing to an HTTP(s) URL or base64 data URL. + +```bash +curl -s http://localhost:8091/v1/videos \ + -F "prompt=A person singing" \ + -F 'image_reference={"image_url": "https://example.com/face.png"}' \ + -F 'audio_reference={"audio_url": "https://example.com/speech.mp3"}' \ + -F "width=832" \ + -F "height=480" \ + -F "num_inference_steps=40" \ + -F "guidance_scale=4.5" \ + -F "fps=16" +``` + + ### Synchronous Generation ```bash @@ -146,7 +204,12 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos ## Model-Specific Examples -For complete text-to-video and image-to-video walkthroughs, see: +For complete text-to-video, image-to-video, and model-specific video-to-video +walkthroughs, see: - [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md) - [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md) +- [Speech-to-Video](../user_guide/examples/online_serving/speech_to_video.md) + for Wan2.2-S2V audio-driven lip-sync generation +- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md) + for model-specific video-to-video examples and conditioning controls diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 3c58c1181..234446c01 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -145,9 +145,10 @@ The following tables show which models support each feature: | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | |------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:| | **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | +| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ | | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | -| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/docs/user_guide/examples/online_serving/speech_to_video.md b/docs/user_guide/examples/online_serving/speech_to_video.md new file mode 100644 index 000000000..80788258c --- /dev/null +++ b/docs/user_guide/examples/online_serving/speech_to_video.md @@ -0,0 +1,128 @@ +# Speech-To-Video + +Source . + +This example demonstrates how to deploy the Wan2.2 speech-to-video (S2V) model for online video generation using vLLM-Omni. + +## Supported Models + +| Model | Model ID | +|-------|----------| +| Wan2.2 S2V (14B) | `Wan-AI/Wan2.2-S2V-14B` | + +## Start Server + +### Basic Start + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +vllm serve Wan-AI/Wan2.2-S2V-14B --omni \ + --model-class-name WanS2VPipeline \ + --tensor-parallel-size 2 \ + --flow-shift 3.0 \ + --vae-use-slicing --vae-use-tiling \ + --cache-backend cache_dit \ + --port 8091 +``` + +### Start with Script + +```bash +bash run_server.sh +``` + +The script allows overriding: +- `MODEL` (default: `Wan-AI/Wan2.2-S2V-14B`) +- `PORT` (default: `8091`) +- `FLOW_SHIFT` (default: `3.0`) +- `TP` (default: `2`) +- `CACHE_BACKEND` (default: `cache_dit`) + +### Start with 4 GPUs + +```bash +TP=4 bash run_server.sh +``` + +## API Endpoints + +The server provides OpenAI-compatible video generation endpoints: + +- `POST /v1/videos`: create a video generation job (async) +- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks) +- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata +- `GET /v1/videos`: list stored video jobs +- `GET /v1/videos/{video_id}/content`: download the generated video file +- `DELETE /v1/videos/{video_id}`: delete the job and any stored output + +## Sync API (Recommended for Testing) + +`POST /v1/videos/sync` blocks until generation completes and returns the raw +video bytes (`video/mp4`) directly in the response body. + +```bash +bash run_curl_speech_to_video.sh +``` + +## Async Job API + +`POST /v1/videos` creates an asynchronous job and returns immediately. + +```bash +no_proxy=127.0.0.1 \ +create_response=$(curl -sS -X POST http://127.0.0.1:8091/v1/videos \ + -H "Accept: application/json" \ + -F "prompt=A person singing" \ + -F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \ + -F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \ + -F "width=832" -F "height=480" \ + -F "num_inference_steps=40" \ + -F "guidance_scale=4.5" \ + -F "fps=16") + +video_id=$(echo "$create_response" | jq -r '.id') + +# Poll until complete +while true; do + status=$(curl -s "http://127.0.0.1:8091/v1/videos/${video_id}" | jq -r '.status') + if [ "$status" = "completed" ]; then + break + fi + if [ "$status" = "failed" ]; then + echo "Video generation failed" + exit 1 + fi + sleep 2 +done + +# Download +curl -L "http://127.0.0.1:8091/v1/videos/${video_id}/content" -o s2v_output.mp4 +``` + +## Request Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `prompt` | str | - | Text description of the desired video | +| `image_reference` | JSON str | - | Image reference: `{"image_url": "..."}` — supports HTTP(s) URLs or base64 data URLs | +| `audio_reference` | JSON str | - | Audio reference: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs | +| `width` | int | None | Video width in pixels | +| `height` | int | None | Video height in pixels | +| `fps` | int | None | Output video frame rate | +| `num_inference_steps` | int | None | Number of denoising steps (40 recommended) | +| `guidance_scale` | float | None | CFG guidance scale | +| `seed` | int | None | Random seed for reproducibility | + +## Notes + +- S2V requires both a reference image (`image_reference`) and an audio reference + (`audio_reference`). The generated video will show a person matching the reference + image with lip movements synchronized to the audio. +- `audio_reference` accepts a JSON string: `{"audio_url": "..."}` where the URL can be: + - An HTTP/HTTPS URL (e.g., `https://example.com/audio.mp3`) + - A base64 data URL (e.g., `data:audio/mp3;base64,...`) +- `--model-class-name WanS2VPipeline` is required on the server to select the + S2V pipeline (distinct from the T2V/I2V pipelines). +- `--cache-backend cache_dit` enables DiT caching for ~2x speedup on cached steps. +- Audio is muxed into the output MP4 automatically. +- Always pass `fps=16` to ensure correct video/audio alignment. diff --git a/docs/user_guide/fault_injection_reliability_matrix.md b/docs/user_guide/fault_injection_reliability_matrix.md new file mode 100644 index 000000000..2b5fd2401 --- /dev/null +++ b/docs/user_guide/fault_injection_reliability_matrix.md @@ -0,0 +1,117 @@ +# vLLM-Omni Failure Mode Injection Scenarios and Expected Behavior Matrix + +## Background and Goals + +The online serving pipeline of vLLM-Omni can be abstracted as: `Serve` (service entry process) -> `Engine` (model orchestration unit) -> `Worker` (execution process) -> `GPU` (compute resource). + +- `Serve` provides external APIs, accepts requests, and manages lifecycle. +- `Engine` handles request scheduling and inference orchestration; the number and organization of engines can vary by model type. +- `Worker` performs concrete model computation, usually bound to underlying GPU resources. +- `GPU` is the resource layer that ultimately carries inference compute and memory usage. + +
+ +```mermaid +%%{init: {'flowchart': {'useMaxWidth': false}}}%% +flowchart TD + R(request)--> A(serve) + A --> B(engine) + A --> C(engine) + A --> D(engine) + B --> E(worker) + C --> F(worker) + D --> I(worker) + E --> G(GPU) + F --> H(GPU) + I --> L(GPU) +``` + +
+ +The diagram above shows an abstract relationship. Omni models and Diffusion models share the same hierarchy; differences mainly come from engine decomposition, worker count, and scheduling strategy. + +In real customer environments, the system will inevitably encounter accidental operations or environmental disturbances, such as unintended process termination, container/node jitter, or short-term resource contention. From the system perspective, many of these events map to two fault categories: "process receives abnormal signals" or "GPU memory is squeezed." For example, pressing `Ctrl+C` essentially sends a `SIGINT` signal to the `serve` process; running `kill -15 ` maps to graceful termination via `SIGTERM`; running `kill -9 ` maps to forced termination via `SIGKILL`; when `docker kill` is executed or a Pod is deleted, the container runtime usually sends `SIGKILL`. + +This document currently covers two core failure modes, and the system is expected to show predictable and observable behavior under each mode: + +- **Failure Mode 1: Process receives abnormal signals (`SIGINT` / `SIGTERM` / `SIGKILL`)** + Focus on three aspects: whether processes exit as expected and complete cleanup, whether requests fail fast or connections are interrupted, and whether GPU resources are eventually fully released with no residue. +- **Failure Mode 2: OOM (occupy all free GPU memory via an extra process)** + Focus on two aspects: whether service health degrades to the expected state (for example, `503`), and whether requests fail within an acceptable time without hanging. + +More failure modes (for example, network jitter and network interruption) will be added in future iterations to improve end-to-end reliability validation coverage. + +> Note: independent fault injection on the `engine` component is not covered in the current version and will be added gradually in future versions. + +## Fault Injection Scenario Matrix + +| Scenario | Fault Type | System Behavior | Current Status | +|------|----------------------------|----------|----------| +| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | | +| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | | +| No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725)
[#43060](https://github.com/vllm-project/vllm/issues/43060) | +| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released | | +| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released | | +| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly | | +| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed | | +| Under load | Send `SIGKILL` to Worker | In-flight requests are hard interrupted (5xx/connection drop); main service becomes unavailable | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | +| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable | | +| Under load | Send `SIGKILL` to serve main process | serve is hard-killed and current connections are interrupted; in-flight requests fail; after cleanup there are no residual processes and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | +| Under load | Send `SIGTERM` to serve main process | serve stops accepting new requests and executes shutdown flow; in-flight requests fail; no residue remains and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | +| Under load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | `Ctrl+C`-style serve shutdown; in-flight requests fail (5xx/connection interruption); service unavailable; after exit there is no residue and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | +| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly | | +| Under load | Send `SIGTERM` to all related processes | All processes exit gracefully; in-flight requests fail; service unavailable; GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | +| OOM | Occupy all free GPU memory via an extra process | After OOM injection process starts, GPU memory is continuously saturated; service enters unavailable/degraded state and health check drops to 503; different request types (chat/speech, etc.) fail fast within a fixed time and return 500 (no hanging) | [#4285](https://github.com/vllm-project/vllm-omni/issues/4285) | + +## Source of Conclusions + +The behaviors and conclusions above are summarized from current fault injection validation results on `Qwen3-Omni` and `Wan2.2`. + +### Example: `SIGTERM` fault injection log (Qwen3-Omni) + +The following full log is from a real run where a `SIGTERM` signal was sent to the serve root process. + +```text +Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml +Server ready on 127.0.0.1:60675 (OmniServer startup took 206.023s) +OmniServer started successfully +[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml +[reliability][process-kill] current_server_proc pid=557118 name=python3 cmdline=/workspace/.venv/bin/python3 -c from multiprocessing.resource_tracker import main;main(57) +[reliability][process-kill] current_server_proc pid=557119 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0 +[reliability][process-kill] current_server_proc pid=557122 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0 +[reliability][process-kill] current_server_proc pid=558201 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0 + +[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml +[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down +[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [async_omni_engine.py:2397] [AsyncOmniEngine] Shutting down Orchestrator +[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:351] [Orchestrator] Received shutdown signal +[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:1359] [Orchestrator] Shutting down all 3 client(s) + +[0;36m(StageEngineCoreProc pid=557119)[0;0m [32mINFO[0m [90m05-19 12:24:51[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0) +[0;36m(StageEngineCoreProc pid=557119)[0;0m [32mINFO[0m [90m05-19 12:24:51[0m [90m[core.py:1265][0m Shutdown complete +[0;36m(APIServer pid=556855)[0;0m INFO: Shutting down +[0;36m(APIServer pid=556855)[0;0m INFO: Waiting for application shutdown. +[0;36m(APIServer pid=556855)[0;0m INFO: Application shutdown complete. +[0;36m(APIServer pid=556855)[0;0m INFO: Finished server process [556855] +[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down +[rank0]:[W519 12:24:53.427894833 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator()) +[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:56 [stage_pool.py:731] [StagePool] Stage 0 replica 0 shut down +[0;36m(StageEngineCoreProc pid=557122)[0;0m [32mINFO[0m [90m05-19 12:24:56[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0) +[0;36m(StageEngineCoreProc pid=557122)[0;0m [32mINFO[0m [90m05-19 12:24:56[0m [90m[core.py:1265][0m Shutdown complete +[rank0]:[W519 12:24:57.569425065 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator()) +[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:25:00 [stage_pool.py:731] [StagePool] Stage 1 replica 0 shut down +[0;36m(StageEngineCoreProc pid=558201)[0;0m [32mINFO[0m [90m05-19 12:25:00[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0) +[0;36m(StageEngineCoreProc pid=558201)[0;0m [32mINFO[0m [90m05-19 12:25:00[0m [90m[core.py:1265][0m Shutdown complete +[0;36m(APIServer pid=556855)[0;0m WARNING 05-19 12:25:01 [async_omni_engine.py:2406] [AsyncOmniEngine] Orchestrator thread did not exit in time +[rank0]:[W519 12:25:01.867242016 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator()) +``` + +### How this maps to the matrix + +- This log corresponds to the `Send SIGTERM to serve main process` scenario in the matrix. +- The `root-kill ... signal=SIGTERM` line confirms the fault injection action is applied as expected. +- The sequence from `Shutting down Orchestrator` to stage-level `Shutdown complete` and `Stage ... shut down` matches the expected graceful shutdown path and child-process cleanup behavior described in the matrix. + +## Related Documentation + +- For practical CI debugging procedures and triage guidance, see [CI Failures](../contributing/ci/failures.md). diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index c8f14548c..1abcd4389 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -28,6 +28,50 @@ in deep DiT blocks. Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. FP8 on Ampere may use a weight-only path where available. +### Faster FP8 GEMM on Blackwell (quack) + +On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which +applies the bias as a separate kernel after the GEMM. On the small GEMMs in video +DiTs this bias add is a significant overhead. Installing the optional `quack` kernel +lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering +that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster). + +```bash +# CUDA 12.9 +pip install vllm-omni[quack] + +# CUDA 13.x +pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130 +``` + +It is enabled automatically once installed (no flag needed) and is **Blackwell-only**: +on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there. +Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is +not installed, FP8 still works — it just keeps the unfused FlashInfer path. + +#### Compile cache and warmup + +quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer +the first time across all autotuned configs). The compiled `.o` files are cached on +disk and reused on later runs, so this is a one-time cost — **not** per request. + +vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with +`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts. +In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into +the image — so the first cold start does not recompile. The engine's startup dummy +run already exercises the kernels, so with a warm cache the first real request is fast. + +To pre-warm specific shapes (e.g. at image build time): + +```python +from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8 +# (M, K, N) per linear; M = number of tokens for your resolution/frame count +warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) +``` + +> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install +> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12. + ## Model Type Support ### Diffusion Model (Qwen-Image, Wan2.2) diff --git a/examples/offline_inference/cosmos3/inputs/v2v.json b/examples/offline_inference/cosmos3/inputs/v2v.json new file mode 100644 index 000000000..c16c4bbb7 --- /dev/null +++ b/examples/offline_inference/cosmos3/inputs/v2v.json @@ -0,0 +1,14 @@ +{ + "prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.", + "negative_prompt": "", + "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4", + "height": 720, + "width": 1280, + "num_frames": 189, + "num_inference_steps": 35, + "guidance_scale": 6.0, + "flow_shift": 10.0, + "fps": 24, + "condition_frame_indexes_vision": [0, 1], + "condition_video_keep": "first" +} diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 7837fe16d..0f2445673 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -143,6 +143,20 @@ def parse_args() -> argparse.Namespace: default=None, required=False, ) + parser.add_argument( + "--width", + type=int, + default=None, + metavar="W", + help="Output image width in pixels. Default: None (pipeline's default).", + ) + parser.add_argument( + "--height", + type=int, + default=None, + metavar="H", + help="Output image height in pixels. Default: None (pipeline's default).", + ) parser.add_argument( "--seed", type=int, @@ -236,8 +250,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--resolution", type=int, - default=640, - help="Bucket in (640, 1024) to determine the condition and output resolution", + default=None, + help="Bucket in (640, 1024) to determine the condition and output resolution. If width and height are not provided, this will be set to default 640.", ) parser.add_argument( @@ -384,6 +398,15 @@ def parse_args() -> argparse.Namespace: def main(): args = parse_args() + if args.resolution and (args.width or args.height): + raise ValueError("--resolution and --width/--height cannot be specified together") + if args.width is not None and args.width <= 0: + raise ValueError("--width must be a positive integer") + if args.height is not None and args.height <= 0: + raise ValueError("--height must be a positive integer") + if not args.width and not args.height and not args.resolution: + args.resolution = 640 + # Validate input images exist and load them input_images = [] for image_path in args.image: @@ -464,6 +487,7 @@ def main(): print( f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}" ) + print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}") print(f"{'=' * 60}\n") generation_start = time.perf_counter() @@ -473,22 +497,26 @@ def main(): omni.start_profile() # Generate edited image + sampling_params = OmniDiffusionSamplingParams( + width=args.width, + height=args.height, + generator=generator, + true_cfg_scale=args.cfg_scale, + guidance_scale=args.guidance_scale, + guidance_scale_2=args.guidance_scale_2, + num_inference_steps=args.num_inference_steps, + num_outputs_per_prompt=args.num_outputs_per_prompt, + layers=args.layers, + ) + if args.resolution: + sampling_params.resolution = args.resolution outputs = omni.generate( { "prompt": args.prompt, "negative_prompt": args.negative_prompt, "multi_modal_data": {"image": input_image}, }, - OmniDiffusionSamplingParams( - generator=generator, - true_cfg_scale=args.cfg_scale, - guidance_scale=args.guidance_scale, - guidance_scale_2=args.guidance_scale_2, - num_inference_steps=args.num_inference_steps, - num_outputs_per_prompt=args.num_outputs_per_prompt, - layers=args.layers, - resolution=args.resolution, - ), + sampling_params, ) generation_end = time.perf_counter() generation_time = generation_end - generation_start diff --git a/examples/offline_inference/speech_to_video/speech_to_video.py b/examples/offline_inference/speech_to_video/speech_to_video.py index 97ce63700..d783ee2f7 100644 --- a/examples/offline_inference/speech_to_video/speech_to_video.py +++ b/examples/offline_inference/speech_to_video/speech_to_video.py @@ -16,9 +16,10 @@ Usage: """ import argparse -import os +import json import time from pathlib import Path +from typing import Any import torch @@ -29,6 +30,16 @@ 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)." @@ -87,6 +98,12 @@ 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, @@ -178,30 +195,10 @@ def parse_args() -> argparse.Namespace: help="Enable cache-dit summary logging after diffusion forward passes.", ) parser.add_argument( - "--profile-dir", - type=str, + "--profiler-config", + type=parse_profiler_config, default=None, - 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.", + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) return parser.parse_args() @@ -239,23 +236,9 @@ def main(): ulysses_degree=args.ulysses_degree, ) - # 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 + profiler_enabled = args.profiler_config is not None - 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( + omni_kwargs = dict( model=args.model, model_class_name="WanS2VPipeline", flow_shift=args.flow_shift, @@ -269,8 +252,12 @@ def main(): cache_backend=args.cache_backend, cache_config=cache_config, enable_cache_dit_summary=args.enable_cache_dit_summary, - profiler_config=profiler_config, + profiler_config=args.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}") @@ -325,13 +312,16 @@ 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, list): + if profile_results and isinstance(profile_results, dict): + traces = profile_results.get("traces", []) print("\n" + "=" * 60) print("PROFILING RESULTS:") - for i, result in enumerate(profile_results): - print(f"\nStage {i}:") - if result: - print(f" {result}") + for rank, trace in enumerate(traces): + print(f"\nRank {rank}:") + if trace: + print(f" • Trace: {trace}") + if not traces: + print(" No traces collected.") print("=" * 60) else: print("[Profiler] No valid profiling data returned.") diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index a0e9b529b..15672b94f 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -17,6 +17,7 @@ list of supported architectures across all modalities, see | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | | GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz | +| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz | | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz | | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz | @@ -159,6 +160,46 @@ Streaming requires `async_chunk: true` in the stage config. --- +## Ming-omni-tts + +Dense 0.5B two-stage TTS pipeline (`AR + flow` + audio VAE) at 44.1 kHz. The example covers style, IP voice, music-only generation, text-to-audio events, emotion, dialect, zero-shot cloning, podcast, speech+BGM, and speech+environment-sound cases. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case style \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +### Voice cloning +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case zero_shot \ + --ref-audio /path/to/reference.wav \ + --ref-text "在此奉劝大家别乱打美白针。" \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +### Streaming +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case basic \ + --ref-audio /path/to/reference.wav \ + --streaming \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +### Notes +- `style`, `ip`, `bgm`, and `tta` do not require reference audio. +- Reference-audio cases use `--ref-audio`; `zero_shot` also requires `--ref-text`. +- `podcast` uses multiple references via `--ref-audio-paths`. +- Full case details live in [`ming_tts/README.md`](ming_tts/README.md). + +--- + ## Ming-flash-omni-TTS Standalone talker-only deployment of Ming-flash-omni-2.0 at 44.1 kHz. Voice is controlled through caption fields (`风格` / `IP` / `语速`/`基频`/`音量`) rather than reference audio. diff --git a/examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py b/examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py new file mode 100644 index 000000000..8b2f3c8a1 --- /dev/null +++ b/examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline higgs-audio v3 inference example. + +Runs Stage 0 (Qwen3 talker) + Stage 1 (HiggsAudio codec) end-to-end +through the vLLM-Omni engine without going through the HTTP server, and +saves a 24 kHz mono WAV per prompt. + +Example (plain TTS): + + python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\ + --texts "Hello world." \\ + --output-dir results/higgs_v3_wavs + +Example (voice clone): + + python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\ + --texts "Hello world." \\ + --ref-audio path/to/reference.wav \\ + --ref-text "Transcript of the reference clip." \\ + --output-dir results/higgs_v3_wavs +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") +os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0") + +import re +import time +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch + +from vllm_omni import Omni + +SAMPLE_RATE = 24_000 +DEFAULT_TEXTS = ( + "Hello world.", + "The quick brown fox jumps over the lazy dog.", + "Today is a beautiful day for a walk in the park.", +) + + +def parse_args(): + import argparse + + parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference") + parser.add_argument( + "--model", + type=str, + default="bosonai/higgs-audio-v3-tts-4b", + help="Stage-0 talker model id or path.", + ) + parser.add_argument( + "--texts", + type=str, + nargs="+", + default=list(DEFAULT_TEXTS), + help="One or more plain-text prompts to synthesize.", + ) + parser.add_argument( + "--output-dir", + type=str, + default="results/higgs_v3_wavs", + help="Directory to write per-prompt WAV files.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path.", + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=2048, + help="Cap on Stage-0 codec frames per request.", + ) + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help="Reference audio for voice clone (WAV/FLAC/MP3 path).", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Transcript of the reference audio. Optional but improves fidelity.", + ) + return parser.parse_args() + + +def _slugify(text: str) -> str: + slug = re.sub(r"\s+", "_", text.strip().lower()) + slug = re.sub(r"[^a-z0-9_]+", "", slug) + return slug[:48] or "prompt" + + +def _extract_pcm(multimodal_output: dict) -> torch.Tensor: + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is None: + raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}") + if isinstance(audio, list): + valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] + if not valid: + raise ValueError("audio list is empty") + return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] + return torch.as_tensor(audio).float().cpu().reshape(-1) + + +def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray: + arr = pcm.numpy() + if arr.dtype.kind == "f": + arr = np.clip(arr, -1.0, 1.0) + arr = (arr * 32767.0).astype(np.int16) + else: + arr = arr.astype(np.int16) + return arr + + +def main(): + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True) + + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + HiggsAudioV3TokenizerAdapter, + apply_delay_pattern, + encode_reference_audio, + ) + + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + adapter = HiggsAudioV3TokenizerAdapter(tokenizer) + + # Load and encode reference audio once (if voice cloning) + ref_codes_delayed: torch.Tensor | None = None + if args.ref_audio is not None: + ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False) + if ref_wav.ndim == 2: + ref_wav = ref_wav.mean(axis=1) + ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr)) + ref_codes_delayed = apply_delay_pattern(ref_codes_raw) + print(f"Reference : {args.ref_audio}") + print(f"Ref codes : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed") + if args.ref_text: + print(f"Ref text : {args.ref_text}") + + print(f"Model : {args.model}") + print(f"Prompts : {len(args.texts)}") + print(f"Output dir : {output_dir}") + print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}") + + total_elapsed = 0.0 + total_dur = 0.0 + for text in args.texts: + if ref_codes_delayed is not None: + prompt_ids = adapter.build_prompt( + text, + num_ref_tokens=int(ref_codes_delayed.shape[0]), + reference_text=args.ref_text, + ) + prompt = { + "prompt_token_ids": prompt_ids, + "additional_information": { + "audio_input_ids": ref_codes_delayed.to(torch.long), + "audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool), + }, + } + else: + prompt_ids = adapter.build_prompt(text) + prompt = {"prompt_token_ids": prompt_ids} + + t_start = time.perf_counter() + outputs = engine.generate([prompt]) + elapsed = time.perf_counter() - t_start + total_elapsed += elapsed + + mm = outputs[0].outputs[0].multimodal_output + pcm = _extract_pcm(mm) + slug = _slugify(text) + out_path = output_dir / f"{slug}.wav" + sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16") + dur = pcm.numel() / SAMPLE_RATE + total_dur += dur + print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}") + + rtf = total_elapsed / total_dur if total_dur > 0 else float("inf") + print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}") + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/text_to_speech/ming_tts/README.md b/examples/offline_inference/text_to_speech/ming_tts/README.md new file mode 100644 index 000000000..3fdfaf4d5 --- /dev/null +++ b/examples/offline_inference/text_to_speech/ming_tts/README.md @@ -0,0 +1,248 @@ +# Ming-omni-tts Offline Inference + +`end2end.py` runs Ming dense 0.5B end to end with vLLM-Omni. It uses the in-repo Ming prompt assembly helper directly, so the example request shape matches the real integration instead of a simplified wrapper. + +## Files + +| File | Purpose | +|---|---| +| `end2end.py` | Driver: CLI, case loading, prompt construction, orchestration (~150 lines) | +| `cases.yaml` | All 11 built-in case definitions (prompt, text, instruction, ref-audio flags, flow controls) | +| `runner.py` | Engine management and audio output (streaming + blocking paths) | + +## Model Overview + +Ming dense 0.5B is exposed here as a two-stage offline pipeline: + +- **Stage 0**: Qwen2-based AR generation with Ming prompt formatting and inline flow controls +- **Stage 1**: audio VAE decode to mono 44.1 kHz waveform + +`config_ming_tts.py` adapts the checkpoint's HuggingFace config fields +(LLM, DiT, aggregator, AudioVAE, token ids). `vllm_omni/deploy/ming_tts.yaml` +selects the vLLM-Omni pipeline and stage runtime topology, including +connectors, async chunking, memory limits, and sampling defaults. + +The example supports both: + +- **Blocking eager** via `vllm_omni/deploy/ming_tts.yaml` +- **Async chunk eager** via `vllm_omni/deploy/ming_tts.yaml` (default `async_chunk: true`) + +## Setup + +Install vLLM-Omni with the platform requirements for your accelerator: + +```bash +uv pip install -e . +``` + +The Ming offline example does not require a separate upstream Ming package. +Reference-audio cases use the repo dependencies for audio loading, +resampling, and CampPlus speaker extraction, including `soundfile`, +`torchaudio`, and `onnxruntime-rocm` in the documented ROCm environment. + +The tested ROCm environment is summarized in the +[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md). + +## Supported Cases + +These cases cover the upstream dense 0.5B cookbook surface that maps cleanly onto the current vLLM-Omni example: + +- `style`: zero-speaker style-conditioned speech +- `ip`: zero-speaker IP voice generation +- `bgm`: music-only generation +- `tta`: text-to-audio event generation with FlowLoss controls +- `emotion`: reference-audio speech with emotion control +- `basic`: reference-audio speech with speed / pitch / volume control +- `dialect`: reference-audio speech with dialect control +- `zero_shot`: reference-audio cloning with explicit transcript +- `podcast`: multi-reference dialogue generation with automatic speaker embedding extraction +- `speech_bgm`: speech with background music conditioning +- `speech_sound`: speech with environmental sound conditioning + +## Quick Start + +Run the zero-speaker style example: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case style \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +Run zero-shot cloning with a transcript: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case zero_shot \ + --ref-audio /path/to/10002287-00000094.wav \ + --ref-text "在此奉劝大家别乱打美白针。" \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +Run emotion-controlled speech: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case emotion \ + --ref-audio /path/to/emotion_prompt.wav \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +Run podcast generation with two reference clips: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case podcast \ + --ref-audio-paths /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +The script automatically extracts one 192-d speaker embedding per reference WAV using the Ming model's `campplus.onnx`. + +If you already have precomputed multi-speaker embeddings, you can override extraction with: + +```bash +--speaker-embedding /path/to/podcast_speaker_embeddings.json +``` + +where the JSON is a list of speaker embeddings, one 192-d vector per speaker. + +Run text-to-audio event generation: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case tta \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +Use async_chunk streaming: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case basic \ + --ref-audio /path/to/10002287-00000095.wav \ + --streaming \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently +supports one prompt per process invocation; use blocking mode for +`--num-prompts > 1`. + +Collect runtime stats and a manifest: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case style \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager \ + --enable-stats \ + --stats-log-file output_audio/ming_style_pipeline.log \ + --metadata-json output_audio/ming_style_manifest.json +``` + +## Reference Fixtures + +The upstream Ming cookbook uses these public audio fixtures from `inclusionAI/Ming-omni-tts/data/wavs`: + +- `10002287-00000094.wav` for zero-shot cloning +- `10002287-00000095.wav` for `basic` +- `emotion_prompt.wav` for `emotion` +- `yue_prompt.wav` for `dialect` +- `00000309-00000300.wav` for `speech_bgm` and `speech_sound` +- `CTS-CN-F2F-2019-11-11-423-012-A.wav` and `CTS-CN-F2F-2019-11-11-423-012-B.wav` for `podcast` + +## Validation Matrix + +The repo-facing example is intended to cover the same dense TTS workflows used +by the local Ming validation script: + +| Case | Blocking `deploy/ming_tts.yaml` | Async chunk `deploy/ming_tts.yaml` | Extra inputs | +|---|---:|---:|---| +| `style` | Yes | Optional smoke test | none | +| `ip` | Yes | Optional smoke test | none | +| `bgm` | Yes | Optional smoke test | none | +| `tta` | Yes | Optional smoke test | none | +| `emotion` | Yes | Yes | `--ref-audio emotion_prompt.wav` | +| `basic` | Yes | Yes | `--ref-audio 10002287-00000095.wav` | +| `dialect` | Yes | Yes | `--ref-audio yue_prompt.wav` | +| `zero_shot` | Yes | Yes | `--ref-audio 10002287-00000094.wav --ref-text ...` | +| `podcast` | Yes | Yes | two `--ref-audio-paths` | +| `speech_bgm` | Yes | Yes | `--ref-audio 00000309-00000300.wav` | +| `speech_sound` | Yes | Yes | `--ref-audio 00000309-00000300.wav` | + +## Validated Outputs + +The following measurements are retained from an earlier L4 CUDA validation; +they are not ROCm benchmark results. Default async_chunk matched blocking +output frame counts and Stage-1 patch counts for every case: + +| Case | Blocking frames / patches / sec | Async chunk frames / patches / sec | +|---|---:|---:| +| `style` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 | +| `ip` | 183456 / 13 / 4.16 | 183456 / 13 / 4.16 | +| `bgm` | 1326528 / 94 / 30.08 | 1326528 / 94 / 30.08 | +| `tta` | 465696 / 33 / 10.56 | 465696 / 33 / 10.56 | +| `emotion` | 324576 / 23 / 7.36 | 324576 / 23 / 7.36 | +| `basic` | 211680 / 15 / 4.80 | 211680 / 15 / 4.80 | +| `dialect` | 239904 / 17 / 5.44 | 239904 / 17 / 5.44 | +| `zero_shot` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 | +| `podcast` | 437472 / 31 / 9.92 | 437472 / 31 / 9.92 | +| `speech_bgm` | 296352 / 21 / 6.72 | 296352 / 21 / 6.72 | +| `speech_sound` | 352800 / 25 / 8.00 | 352800 / 25 / 8.00 | + +## Key Arguments + +| Argument | Description | +|---|---| +| `--model` | Hugging Face repo or local Ming checkpoint path | +| `--deploy-config` | Deploy config YAML. Use `vllm_omni/deploy/ming_tts.yaml` | +| `--case` | Built-in demo case | +| `--ref-audio` | Single reference wav path for cloning-style cases | +| `--ref-audio-paths` | Multiple reference wav paths, used by `podcast` | +| `--ref-text` | Reference transcript. Required for `zero_shot` | +| `--instructions` | Free-form Ming instruction string | +| `--instruction-json` | Structured Ming instruction JSON | +| `--speaker-embedding` | JSON file containing a 192-d speaker embedding | +| `--extract-speaker-embeddings` | Force CampPlus speaker extraction from the provided reference audio paths | +| `--max-decode-steps` | Override `ming_max_decode_steps` | +| `--num-prompts` | Repeat the same case N times. Outputs are indexed when `N > 1` | +| `--streaming` | Use `AsyncOmni` and async_chunk transport | +| `--enforce-eager` | Recommended for Ming dense; non-eager is out of scope | +| `--enable-stats` / `--log-stats` | Enable vLLM-Omni per-request stats logging | +| `--stats-log-file` | Optional path for the stats log | +| `--metadata-json` | Optional path for the run manifest JSON | +| `--stage-init-timeout` | Per-stage initialization timeout in seconds | +| `--init-timeout` | Total initialization timeout in seconds | +| `--batch-timeout` | Batch timeout in seconds | +| `--worker-backend` | `multi_process` or `ray` | +| `--ray-address` | Ray cluster address when using `--worker-backend ray` | + +## Output + +- The script writes one mono 44.1 kHz WAV file per run +- Default output directory: `output_audio/` +- Default filename: `ming_.wav` +- When `--num-prompts > 1`, outputs are indexed as `ming__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. diff --git a/examples/offline_inference/text_to_speech/ming_tts/cases.yaml b/examples/offline_inference/text_to_speech/ming_tts/cases.yaml new file mode 100644 index 000000000..d17b2761a --- /dev/null +++ b/examples/offline_inference/text_to_speech/ming_tts/cases.yaml @@ -0,0 +1,110 @@ +style: + prompt: "Please generate speech based on the following description.\n" + text: "我会一直在这里陪着你，直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗？" + instruction: + 风格: >- + 这是一种ASMR耳语，属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语，声音气音成分重。音量极低，紧贴麦克风，语速极慢，旨在制造触发听者颅内快感的声学刺激。 + use_zero_spk_emb: true + max_decode_steps: 200 + +ip: + prompt: "Please generate speech based on the following description.\n" + text: "这款产品的名字，叫变态坑爹牛肉丸。" + instruction: + IP: "灵小甄" + use_zero_spk_emb: true + max_decode_steps: 200 + +bgm: + prompt: "Please generate music based on the following description.\n" + text: " Genre: 电子舞曲. Mood: 自信 / 坚定. Instrument: 架子鼓. Theme: 节日. Duration: 30s." + instruction: null + use_zero_spk_emb: false + max_decode_steps: 400 + +tta: + prompt: "Please generate audio events based on given text.\n" + text: "Thunder and a gentle rain" + instruction: null + use_zero_spk_emb: false + max_decode_steps: 200 + cfg: 4.5 + sigma: 0.3 + temperature: 2.5 + +emotion: + prompt: "Please generate speech based on the following description.\n" + text: "我竟然抢到了陈奕迅的演唱会门票！太棒了！终于可以现场听一听他的歌声了！" + instruction: + 情感: "高兴" + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +basic: + prompt: "Please generate speech based on the following description.\n" + text: "简单地说，这相当于惠普把消费领域市场拱手相让了。" + instruction: + 语速: "快速" + 基频: "中" + 音量: "中" + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +dialect: + prompt: "Please generate speech based on the following description.\n" + text: "我觉得社会企业同个人都有责任" + instruction: + 方言: "广粤话" + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +zero_shot: + prompt: "Please generate speech based on the following description.\n" + text: "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" + instruction: null + requires_ref_audio: true + requires_ref_text: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +podcast: + prompt: "Please generate speech based on the following description.\n" + text: " speaker_1:你可以说一下，就大概说一下，可能虽然我也不知道，我看过那部电影没有。\n speaker_2:就是那个叫什么，变相一节课的嘛。\n speaker_1:嗯。\n speaker_2:一部搞笑的电影。\n speaker_1:一部搞笑的。\n" + instruction: null + prompt_text: " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试，做个，当服务员还要去笔试了\n speaker_2:对啊，这真的很奇怪，就是 单纯的因，单纯自己工资不高，只是因为可能人家那个店比较出名一点，就对你苛刻要求\n" + requires_ref_audio_count: 2 + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +speech_bgm: + prompt: "Please generate speech based on the following description.\n" + text: "此次业绩下滑原因，可归结为企业停止服务某些品牌，而带来的负面影响。" + instruction: + BGM: + Genre: "当代古典音乐." + Mood: "温暖 / 友善." + Instrument: "电吉他" + Theme: "节日." + SNR: 10.0 + ENV: null + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +speech_sound: + prompt: "Please generate speech based on the following description.\n" + text: "此次业绩下滑原因，可归结为企业停止服务某些品牌，而带来的负面影响。" + instruction: + BGM: + ENV: "Birds chirping" + SNR: 10.0 + Genre: null + Mood: null + Instrument: null + Theme: null + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 diff --git a/examples/offline_inference/text_to_speech/ming_tts/end2end.py b/examples/offline_inference/text_to_speech/ming_tts/end2end.py new file mode 100644 index 000000000..1a6991e40 --- /dev/null +++ b/examples/offline_inference/text_to_speech/ming_tts/end2end.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline inference demo for Ming-omni-tts via vLLM Omni.""" + +import json +import os +from pathlib import Path + +import soundfile as sf +import torch +import torchaudio +import yaml +from transformers import AutoTokenizer +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + KEY_CFG, + KEY_MAX_DECODE_STEPS, + KEY_SIGMA, + KEY_TEMPERATURE, + SAMPLE_RATE, +) +from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt +from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor + +try: + from .runner import ( + build_manifest, + build_sampling_params, + resolve_metadata_json, + resolve_stats_log_file, + run_generation, + ) +except ImportError: + from runner import ( + build_manifest, + build_sampling_params, + resolve_metadata_json, + resolve_stats_log_file, + run_generation, + ) + +_DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B" +_REPO_ROOT = Path(__file__).resolve().parents[4] +_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni/deploy/ming_tts.yaml") +_CASES_FILE = Path(__file__).with_name("cases.yaml") + +CASE_DEFAULTS = yaml.safe_load(_CASES_FILE.read_text(encoding="utf-8")) or {} +if not CASE_DEFAULTS: + raise RuntimeError(f"Empty or missing case definitions in {_CASES_FILE}") + + +def _build_parser() -> FlexibleArgumentParser: + p = FlexibleArgumentParser(description="Offline Ming-omni-tts example") + p.add_argument("--model", default=_DEFAULT_MODEL, help="Model name or local path") + p.add_argument("--deploy-config", default=None, help="Deploy config YAML; auto-selected when omitted") + p.add_argument("--case", choices=sorted(CASE_DEFAULTS), default="style", help="Built-in demo case") + p.add_argument("--text", default=None, help="Override case text") + p.add_argument("--prompt", default=None, help="Override the system prompt prefix") + p.add_argument("--instructions", default=None, help="Free-form Ming instruction string") + p.add_argument( + "--instruction-json", default=None, help='Structured Ming instruction JSON, e.g. \'{"方言":"广粤话"}\'' + ) + p.add_argument("--ref-audio", default=None, help="Single reference audio path for cloning-style cases") + p.add_argument("--ref-audio-paths", nargs="+", default=None, help="Multiple reference audio paths (podcast)") + p.add_argument("--ref-text", default=None, help="Reference transcript; required for zero_shot") + p.add_argument("--speaker-embedding", default=None, help="Path to a JSON speaker embedding file") + p.add_argument( + "--extract-speaker-embeddings", + action="store_true", + help="Extract speaker embeddings from ref audio via campplus.onnx", + ) + p.add_argument("--max-decode-steps", type=int, default=None, help="Override ming_max_decode_steps") + p.add_argument("--output-dir", default="output_audio", help="Directory for output wav files") + p.add_argument("--output-name", default=None, help="Output wav filename") + p.add_argument("--num-prompts", type=int, default=1, help="Repeat the same prompt N times") + p.add_argument("--streaming", action="store_true", help="Use AsyncOmni with async_chunk streaming") + p.add_argument("--trust-remote-code", action="store_true") + p.add_argument("--enforce-eager", action="store_true") + p.add_argument( + "--log-stats", "--enable-stats", dest="log_stats", action="store_true", help="Enable Omni stats logging" + ) + p.add_argument("--stats-log-file", default=None, help="Path for the Omni stats log file") + p.add_argument("--metadata-json", default=None, help="Path for the run manifest JSON") + p.add_argument("--stage-init-timeout", type=int, default=300, help="Per-stage init timeout (s)") + p.add_argument("--init-timeout", type=int, default=600, help="Total init timeout (s)") + p.add_argument("--batch-timeout", type=int, default=5, help="Batch timeout (s)") + p.add_argument("--shm-threshold-bytes", type=int, default=65536) + p.add_argument("--worker-backend", default="multi_process", choices=["multi_process", "ray"]) + p.add_argument("--ray-address", default=None, help="Ray cluster address (--worker-backend ray)") + return p + + +def _finalize_args(args) -> None: + if args.instructions is not None and args.instruction_json is not None: + raise RuntimeError("Use either --instructions or --instruction-json, not both") + if args.num_prompts < 1: + raise RuntimeError("--num-prompts must be at least 1") + if args.streaming and args.num_prompts != 1: + raise RuntimeError("--streaming currently supports exactly one prompt") + if args.deploy_config is None: + args.deploy_config = _DEFAULT_DEPLOY_CONFIG + + +def _load_waveform(path: str) -> torch.Tensor: + samples, sr = sf.read(path, dtype="float32") + wav = torch.as_tensor(samples, dtype=torch.float32) + if wav.ndim == 2: + wav = wav.mean(dim=1) + wav = wav.reshape(1, -1) + if int(sr) != SAMPLE_RATE: + wav = torchaudio.functional.resample(wav, int(sr), SAMPLE_RATE) + return wav + + +def _load_speaker_embedding(path: str) -> torch.Tensor: + return torch.as_tensor(json.loads(Path(path).read_text(encoding="utf-8")), dtype=torch.float32) + + +def _ref_audio_paths(args) -> list[str]: + if args.ref_audio is not None and args.ref_audio_paths is not None: + raise RuntimeError("Use either --ref-audio or --ref-audio-paths, not both") + if args.ref_audio_paths is not None: + return list(args.ref_audio_paths) + return [args.ref_audio] if args.ref_audio else [] + + +def _resolve_reference_inputs(args, case: dict, paths: list[str]): + required = int(case.get("requires_ref_audio_count", 0)) + if required > 0 and len(paths) < required: + raise RuntimeError(f"Case '{args.case}' needs {required} ref audio paths via --ref-audio-paths") + if required <= 0 and case.get("requires_ref_audio") and not paths: + raise RuntimeError(f"--ref-audio required for case '{args.case}'") + if not paths: + return None + wavs = [_load_waveform(p) for p in paths] + return wavs[0] if len(wavs) == 1 else wavs + + +def _resolve_speaker_embedding(args, case: dict, paths: list[str]): + if args.speaker_embedding: + return _load_speaker_embedding(args.speaker_embedding) + if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths: + return None + embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths) + if not embs: + raise RuntimeError("Speaker extraction produced no embeddings") + return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0) + + +def _build_prompt_payload(tokenizer, args): + case = CASE_DEFAULTS[args.case] + paths = _ref_audio_paths(args) + prompt_text = args.ref_text if args.ref_text is not None else case.get("prompt_text") + if case.get("requires_ref_text") and not prompt_text: + raise RuntimeError(f"--ref-text required for case '{args.case}'") + prompt_waveform = _resolve_reference_inputs(args, case, paths) if prompt_text is not None else None + speaker_embedding = _resolve_speaker_embedding(args, case, paths) + use_zero_spk_emb = bool(case.get("use_zero_spk_emb")) and prompt_waveform is None and speaker_embedding is None + runtime_controls = {KEY_MAX_DECODE_STEPS: args.max_decode_steps or case["max_decode_steps"]} + for key, field in [(KEY_CFG, "cfg"), (KEY_SIGMA, "sigma"), (KEY_TEMPERATURE, "temperature")]: + if field in case: + runtime_controls[key] = case[field] + instruction = ( + json.loads(args.instruction_json) if args.instruction_json else (args.instructions or case.get("instruction")) + ) + return build_ming_dense_prompt( + tokenizer, + prompt=args.prompt or case["prompt"], + text=args.text or case["text"], + runtime_controls=runtime_controls, + instruction=instruction, + prompt_text=prompt_text, + prompt_waveform=prompt_waveform, + speaker_embedding=speaker_embedding, + use_zero_spk_emb=use_zero_spk_emb, + ) + + +def main(): + args = _build_parser().parse_args() + _finalize_args(args) + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False) + prompt_payload = _build_prompt_payload(tokenizer, args) + case = CASE_DEFAULTS[args.case] + sampling_params_list = build_sampling_params(args.max_decode_steps or case["max_decode_steps"]) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + stats_log_file = resolve_stats_log_file(args) + summaries = run_generation(args, prompt_payload, sampling_params_list, output_dir, stats_log_file) + metadata_json = resolve_metadata_json(args) + manifest = build_manifest(args, prompt_payload, stats_log_file, summaries) + if metadata_json is not None: + Path(metadata_json).write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8") + print(f"Saved run manifest to {metadata_json}") + + +if __name__ == "__main__": + os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + main() diff --git a/examples/offline_inference/text_to_speech/ming_tts/runner.py b/examples/offline_inference/text_to_speech/ming_tts/runner.py new file mode 100644 index 000000000..35852ce04 --- /dev/null +++ b/examples/offline_inference/text_to_speech/ming_tts/runner.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import json +import time +import uuid +import wave +from pathlib import Path + +import torch +from vllm import SamplingParams + +from vllm_omni import AsyncOmni, Omni +from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + KEY_SPEAKER_EMBEDDING, + SAMPLE_RATE, + TEXT_EOS_TOKEN_ID, +) + + +def coerce_audio_tensor(audio, *, async_chunk: bool) -> torch.Tensor: + if isinstance(audio, list): + parts = [] + for item in audio: + tensor = torch.as_tensor(item, dtype=torch.float32).reshape(-1) + if tensor.numel() > 0: + parts.append(tensor) + if not parts: + return torch.zeros((0,), dtype=torch.float32) + return torch.cat(parts, dim=0) + + return torch.as_tensor(audio, dtype=torch.float32).reshape(-1) + + +def resolve_sr(sr) -> int: + if isinstance(sr, list): + sr = sr[-1] + if hasattr(sr, "item"): + return int(sr.item()) + return int(sr) + + +def extract_sample_rate(multimodal_output: dict) -> int: + sr = multimodal_output.get("sr") + if sr is None: + raise RuntimeError("Expected multimodal_output['sr']") + return resolve_sr(sr) + + +def write_wav(path: str, audio: torch.Tensor, sample_rate: int) -> None: + audio = audio.clamp(-1.0, 1.0) + pcm16 = (audio * 32767.0).round().to(torch.int16).cpu().numpy() + with wave.open(path, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(int(sample_rate)) + wav_file.writeframes(pcm16.tobytes()) + + +def request_index(request_id: str | None, fallback: int) -> int: + try: + return int(request_id) + except (TypeError, ValueError): + if isinstance(request_id, str): + head = request_id.split("_", 1)[0] + if head.isdigit(): + return int(head) + return fallback + + +def audio_summary(audio: torch.Tensor, sample_rate: int) -> dict: + waveform = audio.detach().cpu().reshape(-1).to(torch.float32) + return { + "sample_rate": int(sample_rate), + "num_samples": int(waveform.numel()), + "duration_seconds": float(waveform.numel()) / float(sample_rate), + "max_abs_amplitude": float(waveform.abs().max().item()) if waveform.numel() > 0 else 0.0, + } + + +def resolve_output_name(output_name: str | None, case: str, index: int, total: int) -> str: + if total == 1: + return output_name or f"ming_{case}.wav" + base = Path(output_name or f"ming_{case}.wav") + return f"{base.stem}_{index:05d}{base.suffix or '.wav'}" + + +def resolve_stats_log_file(args) -> str | None: + if not args.log_stats: + return None + if args.stats_log_file: + return args.stats_log_file + base = Path(args.output_name or f"ming_{args.case}.wav").stem + return str(Path(args.output_dir) / f"{base}_pipeline.log") + + +def resolve_metadata_json(args) -> str | None: + if args.metadata_json: + return args.metadata_json + if args.log_stats: + base = Path(args.output_name or f"ming_{args.case}.wav").stem + return str(Path(args.output_dir) / f"{base}_manifest.json") + return None + + +def build_manifest(args, prompt_payload, stats_log_file: str | None, outputs: list[dict]) -> dict: + additional_information = {} + if isinstance(prompt_payload, dict): + additional_information = dict(prompt_payload.get("additional_information", {})) + return { + "model": args.model, + "case": args.case, + "streaming": bool(args.streaming), + "deploy_config": args.deploy_config, + "enforce_eager": bool(args.enforce_eager), + "num_prompts": int(args.num_prompts), + "log_stats": bool(args.log_stats), + "stats_log_file": stats_log_file, + "prompt_text": additional_information.get("prompt_text"), + "instruction": additional_information.get("instruction"), + "speaker_embedding_shape": ( + list(additional_information[KEY_SPEAKER_EMBEDDING].shape) + if KEY_SPEAKER_EMBEDDING in additional_information + and hasattr(additional_information[KEY_SPEAKER_EMBEDDING], "shape") + else None + ), + "outputs": outputs, + "generated_at_unix": time.time(), + } + + +def build_engine_kwargs(args, stats_log_file: str | None) -> dict: + kwargs = { + "model": args.model, + "deploy_config": args.deploy_config, + "enforce_eager": args.enforce_eager, + "trust_remote_code": args.trust_remote_code, + "log_stats": args.log_stats, + "stage_init_timeout": args.stage_init_timeout, + "init_timeout": args.init_timeout, + "batch_timeout": args.batch_timeout, + "shm_threshold_bytes": args.shm_threshold_bytes, + "worker_backend": args.worker_backend, + } + if stats_log_file is not None: + kwargs["log_file"] = stats_log_file + if args.ray_address is not None: + kwargs["ray_address"] = args.ray_address + return kwargs + + +def build_sampling_params(max_decode_steps: int) -> list[SamplingParams]: + return [ + SamplingParams( + temperature=0.0, + max_tokens=max_decode_steps + 1, + stop_token_ids=[int(TEXT_EOS_TOKEN_ID)], + ), + SamplingParams(temperature=0.0, max_tokens=1), + ] + + +async def run_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None): + engine = AsyncOmni(**build_engine_kwargs(args, stats_log_file)) + try: + all_audio_chunks = [] + accumulated_samples = 0 + chunk_idx = 0 + start_time = time.time() + chunk_times = [] + ttfp_seconds = None + final_stage_output = None + async for stage_output in engine.generate( + prompt=prompt_payload, + request_id=str(uuid.uuid4()), + sampling_params_list=sampling_params_list, + ): + final_stage_output = stage_output + multimodal_output = stage_output.multimodal_output or {} + audio = multimodal_output.get("audio") + if audio is None: + continue + + finished = stage_output.finished + if isinstance(audio, torch.Tensor): + if finished: + audio_chunk = audio[accumulated_samples:].float().detach().cpu() + else: + audio_chunk = audio.float().detach().cpu() + elif isinstance(audio, list): + audio_chunk = torch.as_tensor(audio[chunk_idx], dtype=torch.float32).reshape(-1).cpu() + else: + audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu() + + accumulated_samples += int(audio_chunk.numel()) + chunk_idx += 1 + if audio_chunk.numel() > 0: + now = time.time() + if ttfp_seconds is None: + ttfp_seconds = now - start_time + chunk_times.append(now) + all_audio_chunks.append(audio_chunk) + + if not all_audio_chunks: + raise RuntimeError("Streaming Ming example produced no audio chunks") + + waveform = torch.cat(all_audio_chunks, dim=0) + output_name = resolve_output_name(args.output_name, args.case, 0, 1) + output_path = str(output_dir / output_name) + write_wav(output_path, waveform, SAMPLE_RATE) + summary = { + "request_id": getattr(final_stage_output, "request_id", None), + "stage_id": getattr(final_stage_output, "stage_id", None), + "output_path": output_path, + "stage_durations": getattr(final_stage_output, "stage_durations", {}), + "peak_memory_mb": getattr(final_stage_output, "peak_memory_mb", 0.0), + "ttfp_seconds": ttfp_seconds, + "mean_inter_chunk_seconds": ( + sum(t1 - t0 for t0, t1 in zip(chunk_times, chunk_times[1:])) / (len(chunk_times) - 1) + if len(chunk_times) > 1 + else None + ), + } + summary.update(audio_summary(waveform, SAMPLE_RATE)) + print(f"Saved streaming output to {output_path}") + print(json.dumps(summary, ensure_ascii=False, indent=2)) + return [summary] + finally: + engine.shutdown() + + +def run_non_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None): + engine = Omni(**build_engine_kwargs(args, stats_log_file)) + try: + outputs = engine.generate( + prompts=[prompt_payload for _ in range(args.num_prompts)], + sampling_params_list=sampling_params_list, + py_generator=False, + ) + summaries = [] + for fallback_index, output in enumerate(outputs): + if output.final_output_type != "audio": + continue + multimodal_output = output.multimodal_output or {} + waveform = coerce_audio_tensor(multimodal_output.get("audio"), async_chunk=False) + sample_rate = extract_sample_rate(multimodal_output) + output_name = resolve_output_name( + args.output_name, + args.case, + request_index(output.request_id, fallback_index), + args.num_prompts, + ) + output_path = str(output_dir / output_name) + write_wav(output_path, waveform, sample_rate) + summary = { + "request_id": output.request_id, + "stage_id": output.stage_id, + "output_path": output_path, + "stage_durations": output.stage_durations, + "peak_memory_mb": output.peak_memory_mb, + } + summary.update(audio_summary(waveform, sample_rate)) + summaries.append(summary) + print(f"Saved output to {output_path}") + print(json.dumps(summary, ensure_ascii=False, indent=2)) + if not summaries: + raise RuntimeError("Non-streaming Ming example produced no audio outputs") + return summaries + finally: + engine.close() + + +def run_generation(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None): + if args.streaming: + return asyncio.run(run_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)) + return run_non_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file) diff --git a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py index c740197b5..c9225c212 100644 --- a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py @@ -212,7 +212,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir output_audio_dur = 0.0 for batch_idx, o in enumerate(outputs): - audio_tensor = torch.cat(o.multimodal_output["audio"]) + audio_tensor = torch.cat((o.multimodal_output["audio"],)) audio_array = audio_tensor.tolist() output_audio_dur += float(len(audio_array)) / 24000 if args.write_audio: diff --git a/examples/online_serving/speech_to_video/README.md b/examples/online_serving/speech_to_video/README.md new file mode 100644 index 000000000..b074a63ce --- /dev/null +++ b/examples/online_serving/speech_to_video/README.md @@ -0,0 +1,115 @@ +# Speech-To-Video + +This example demonstrates how to deploy the Wan2.2 speech-to-video (S2V) model for online video generation using vLLM-Omni. + +## Supported Models + +| Model | Model ID | +|-------|----------| +| Wan2.2 S2V (14B) | `Wan-AI/Wan2.2-S2V-14B` | + +## Start Server + +### Basic Start + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +vllm serve Wan-AI/Wan2.2-S2V-14B --omni \ + --model-class-name WanS2VPipeline \ + --tensor-parallel-size 2 \ + --flow-shift 3.0 \ + --vae-use-slicing --vae-use-tiling \ + --cache-backend cache_dit \ + --port 8091 +``` + +### Start with Script + +```bash +bash run_server.sh +``` + +The script allows overriding: +- `MODEL` (default: `Wan-AI/Wan2.2-S2V-14B`) +- `PORT` (default: `8091`) +- `FLOW_SHIFT` (default: `3.0`) +- `TP` (default: `2`) +- `CACHE_BACKEND` (default: `cache_dit`) + +### Start with 4 GPUs + +```bash +TP=4 bash run_server.sh +``` + +## Sync API (Recommended for Testing) + +`POST /v1/videos/sync` blocks until generation completes and returns the raw +video bytes (`video/mp4`) directly in the response body. + +```bash +bash run_curl_speech_to_video.sh +``` + +## Async Job API + +`POST /v1/videos` creates an asynchronous job and returns immediately. + +```bash +no_proxy=127.0.0.1 \ +create_response=$(curl -sS -X POST http://127.0.0.1:8091/v1/videos \ + -H "Accept: application/json" \ + -F "prompt=A person singing" \ + -F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \ + -F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \ + -F "width=832" -F "height=480" \ + -F "num_inference_steps=40" \ + -F "guidance_scale=4.5" \ + -F "fps=16") + +video_id=$(echo "$create_response" | jq -r '.id') + +# Poll until complete +while true; do + status=$(curl -s "http://127.0.0.1:8091/v1/videos/${video_id}" | jq -r '.status') + if [ "$status" = "completed" ]; then + break + fi + if [ "$status" = "failed" ]; then + echo "Video generation failed" + exit 1 + fi + sleep 2 +done + +# Download +curl -L "http://127.0.0.1:8091/v1/videos/${video_id}/content" -o s2v_output.mp4 +``` + +## Request Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `prompt` | str | - | Text description of the desired video | +| `image_reference` | JSON str | - | Image reference: `{"image_url": "..."}` — supports HTTP(s) URLs or base64 data URLs | +| `audio_reference` | JSON str | - | Audio reference: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs | +| `width` | int | None | Video width in pixels | +| `height` | int | None | Video height in pixels | +| `fps` | int | None | Output video frame rate | +| `num_inference_steps` | int | None | Number of denoising steps (40 recommended) | +| `guidance_scale` | float | None | CFG guidance scale | +| `seed` | int | None | Random seed for reproducibility | + +## Notes + +- S2V requires both a reference image (`image_reference`) and an audio reference + (`audio_reference`). The generated video will show a person matching the reference + image with lip movements synchronized to the audio. +- `audio_reference` accepts a JSON string: `{"audio_url": "..."}` where the URL can be: + - An HTTP/HTTPS URL (e.g., `https://example.com/audio.mp3`) + - A base64 data URL (e.g., `data:audio/mp3;base64,...`) +- `--model-class-name WanS2VPipeline` is required on the server to select the + S2V pipeline (distinct from the T2V/I2V pipelines). +- `--cache-backend cache_dit` enables DiT caching for ~2x speedup on cached steps. +- Audio is muxed into the output MP4 automatically. +- Always pass `fps=16` to ensure correct video/audio alignment. diff --git a/examples/online_serving/speech_to_video/run_curl_speech_to_video.sh b/examples/online_serving/speech_to_video/run_curl_speech_to_video.sh new file mode 100755 index 000000000..dee0d0fac --- /dev/null +++ b/examples/online_serving/speech_to_video/run_curl_speech_to_video.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# Wan2.2 S2V (speech-to-video) curl example using the sync video API. + +set -euo pipefail + +BASE_URL="${BASE_URL:-http://localhost:8091}" +OUTPUT_PATH="${OUTPUT_PATH:-s2v_480p_serve.mp4}" +IMAGE_URL="${IMAGE_URL:-https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png}" +AUDIO_URL="${AUDIO_URL:-https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3}" +PROMPT="${PROMPT:-A person singing}" +WIDTH="${WIDTH:-832}" +HEIGHT="${HEIGHT:-480}" +NUM_INFERENCE_STEPS="${NUM_INFERENCE_STEPS:-40}" +GUIDANCE_SCALE="${GUIDANCE_SCALE:-4.5}" +FPS="${FPS:-16}" + +echo "Sending S2V request..." +echo " Image URL: $IMAGE_URL" +echo " Audio URL: $AUDIO_URL" +echo " Prompt: $PROMPT" +echo " Resolution: ${WIDTH}x${HEIGHT}" +echo " Steps: $NUM_INFERENCE_STEPS" +echo " FPS: $FPS" + +IMAGE_REF_JSON="{\"image_url\": \"${IMAGE_URL}\"}" +AUDIO_REF_JSON="{\"audio_url\": \"${AUDIO_URL}\"}" + +no_proxy=127.0.0.1 \ +curl -X POST "${BASE_URL}/v1/videos/sync" \ + -F "prompt=${PROMPT}" \ + -F "image_reference=${IMAGE_REF_JSON}" \ + -F "audio_reference=${AUDIO_REF_JSON}" \ + -F "width=${WIDTH}" -F "height=${HEIGHT}" \ + -F "num_inference_steps=${NUM_INFERENCE_STEPS}" \ + -F "guidance_scale=${GUIDANCE_SCALE}" \ + -F "fps=${FPS}" \ + --output "${OUTPUT_PATH}" + +if [ -f "$OUTPUT_PATH" ] && [ -s "$OUTPUT_PATH" ]; then + echo "Saved video to ${OUTPUT_PATH} ($(du -h "$OUTPUT_PATH" | cut -f1))" +else + echo "ERROR: Output file is empty or missing" + exit 1 +fi diff --git a/examples/online_serving/speech_to_video/run_server.sh b/examples/online_serving/speech_to_video/run_server.sh new file mode 100755 index 000000000..1f20a5b72 --- /dev/null +++ b/examples/online_serving/speech_to_video/run_server.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# Wan2.2 S2V (speech-to-video) online serving startup script + +MODEL="${MODEL:-Wan-AI/Wan2.2-S2V-14B}" +PORT="${PORT:-8091}" +FLOW_SHIFT="${FLOW_SHIFT:-3.0}" +TP="${TP:-2}" +CACHE_BACKEND="${CACHE_BACKEND:-cache_dit}" + +echo "Starting Wan2.2 S2V server..." +echo "Model: $MODEL" +echo "Port: $PORT" +echo "Flow shift: $FLOW_SHIFT" +echo "Tensor parallel size: $TP" +echo "Cache backend: $CACHE_BACKEND" + +CACHE_BACKEND_FLAG="" +if [ "$CACHE_BACKEND" != "none" ]; then + CACHE_BACKEND_FLAG="--cache-backend $CACHE_BACKEND" +fi + +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +vllm serve "$MODEL" --omni \ + --model-class-name WanS2VPipeline \ + --tensor-parallel-size "$TP" \ + --flow-shift "$FLOW_SHIFT" \ + --vae-use-slicing --vae-use-tiling \ + --port "$PORT" \ + $CACHE_BACKEND_FLAG diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index a3b27b722..c9498596e 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see |---|---|---|---|---|---| | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | +| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — | | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ | | OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — | @@ -202,6 +203,51 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server --- +## Ming-omni-tts + +Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`. + +### Launch +```bash +bash examples/online_serving/text_to_speech/ming_tts/run_server.sh +``` +Equivalent manual command: +```bash +vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --host 0.0.0.0 --port 8091 \ + --enforce-eager --omni +``` + +### Sending requests +```bash +python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ + --text "你好，这是 Ming 在线语音合成测试。" +``` + +Structured dialect control: +```bash +python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ + --text "我觉得社会企业同个人都有责任" \ + --instruction-json '{"方言":"广粤话"}' \ + --ref-audio /path/to/yue_prompt.wav +``` + +Zero-shot cloning: +```bash +python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ + --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \ + --ref-audio /path/to/10002287-00000094.wav \ + --ref-text "在此奉劝大家别乱打美白针。" +``` + +### Notes +- `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook. +- Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples. +- Full request details live in [`ming_tts/README.md`](ming_tts/README.md). + +--- + ## Ming-flash-omni-TTS Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`. diff --git a/examples/online_serving/text_to_speech/fish_speech/gradio_demo.py b/examples/online_serving/text_to_speech/fish_speech/gradio_demo.py index 7db4e0441..cd2279476 100644 --- a/examples/online_serving/text_to_speech/fish_speech/gradio_demo.py +++ b/examples/online_serving/text_to_speech/fish_speech/gradio_demo.py @@ -51,13 +51,18 @@ def build_payload( response_format: str = "wav", stream: bool = False, ) -> dict: - """Build the /v1/audio/speech request payload.""" + """Build the /v1/audio/speech request payload. + + Fish Speech has no built-in speakers; omit ``voice`` so the model + generates with a random timbre (controllable via text tags like + ``[cool]``, ``[warm]``, etc.). When ``ref_audio`` is supplied we + enter voice-cloning mode instead. + """ if not text or not text.strip(): raise gr.Error("Please enter text to synthesize.") payload: dict = { "input": text.strip(), - "voice": "default", "response_format": "pcm" if stream else response_format, "stream": stream, } @@ -193,7 +198,8 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.5): with gr.Accordion("Voice Cloning (optional)", open=False): gr.Markdown( - "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice." + "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice. " + "Leave empty for random timbre; use text tags like `[cool]`, `[warm]`, `[happy]` to steer the output." ) ref_audio = gr.Audio( label="Reference Audio", diff --git a/examples/online_serving/text_to_speech/fish_speech/speech_client.py b/examples/online_serving/text_to_speech/fish_speech/speech_client.py index e3ace37e5..43acffdc7 100644 --- a/examples/online_serving/text_to_speech/fish_speech/speech_client.py +++ b/examples/online_serving/text_to_speech/fish_speech/speech_client.py @@ -39,7 +39,6 @@ def run_tts(args) -> None: payload = { "model": args.model, "input": args.text, - "voice": "default", "response_format": args.response_format, } diff --git a/examples/online_serving/text_to_speech/higgs_audio_v3/README.md b/examples/online_serving/text_to_speech/higgs_audio_v3/README.md new file mode 100644 index 000000000..d1ca8f7b7 --- /dev/null +++ b/examples/online_serving/text_to_speech/higgs_audio_v3/README.md @@ -0,0 +1,40 @@ +# Higgs-Audio V3 Online Serving + +## Start the server + +```bash +# Default: GPU 0, port 8095 +./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh + +# Custom GPU / port +PORT=8096 GPUS=0,1 ./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh +``` + +## Plain text TTS + +```bash +python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \ + --base-url http://localhost:8095 \ + --output-dir /tmp/higgs_v3_batch \ + --prompts "Hello world." "The quick brown fox jumps over the lazy dog." +``` + +## Voice clone + +```bash +python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \ + --base-url http://localhost:8095 \ + --output-dir /tmp/higgs_v3_clone \ + --ref-audio path/to/reference.wav \ + --ref-text "transcript of the reference clip" \ + --prompts "Text to synthesize in the cloned voice." +``` + +## curl example + +```bash +curl -X POST http://localhost:8095/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{"model": "higgs_audio_v3", "input": "Hello world."}' \ + --output hello.wav +``` diff --git a/examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py b/examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py new file mode 100644 index 000000000..0941ba7dc --- /dev/null +++ b/examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Batch client for the higgs-audio v3 online server. + +Sends prompts to ``/v1/audio/speech`` and saves the returned WAV files. + +Usage (plain text -> speech): + + python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \ + --base-url http://localhost:8095 \ + --output-dir /tmp/higgs_v3_batch \ + --prompts "Hello world." "The quick brown fox jumps over the lazy dog." + +Usage (voice clone): + + python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \ + --base-url http://localhost:8095 \ + --output-dir /tmp/higgs_v3_clone \ + --ref-audio path/to/reference.wav \ + --ref-text "the transcript of the reference clip" \ + --prompts "Hello world." +""" + +from __future__ import annotations + +import argparse +import base64 +import sys +from pathlib import Path + +DEFAULT_PROMPTS = ( + "Hello world.", + "The quick brown fox jumps over the lazy dog.", + "Today is a beautiful day for a walk in the park.", + "Innovation distinguishes between a leader and a follower.", +) + + +def _slug(text: str) -> str: + import re + + s = re.sub(r"\s+", "_", text.strip().lower()) + return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--base-url", default="http://localhost:8095") + parser.add_argument("--model", default="higgs_audio_v3") + parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS)) + parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_v3_batch")) + parser.add_argument("--format", choices=("wav", "pcm"), default="wav") + parser.add_argument("--max-new-tokens", type=int, default=2048) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--timeout-s", type=float, default=120.0) + parser.add_argument( + "--ref-audio", + type=Path, + default=None, + help="Reference clip for voice clone (WAV/FLAC/MP3 path). Pair with --ref-text.", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Transcript of the reference clip. Optional but improves fidelity.", + ) + args = parser.parse_args() + + ref_audio_data_url: str | None = None + if args.ref_audio is not None: + if not args.ref_audio.exists(): + print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr) + return 2 + mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg" + ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii") + ref_audio_data_url = f"data:{mime};base64,{ref_b64}" + + try: + import httpx + except ImportError: + print("this client needs `httpx`. Install with `pip install httpx`.", file=sys.stderr) + return 2 + + args.output_dir.mkdir(parents=True, exist_ok=True) + url = args.base_url.rstrip("/") + "/v1/audio/speech" + failures = 0 + with httpx.Client(timeout=args.timeout_s) as client: + for prompt in args.prompts: + payload = { + "model": args.model, + "input": prompt, + "response_format": args.format, + "max_new_tokens": args.max_new_tokens, + "seed": args.seed, + } + if ref_audio_data_url is not None: + payload["ref_audio"] = ref_audio_data_url + if args.ref_text: + payload["ref_text"] = args.ref_text + resp = client.post(url, json=payload) + if resp.status_code != 200: + print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr) + failures += 1 + continue + suffix = ".wav" if args.format == "wav" else ".pcm" + out = args.output_dir / f"{_slug(prompt)}{suffix}" + out.write_bytes(resp.content) + print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)") + + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh b/examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh new file mode 100755 index 000000000..0694a30e7 --- /dev/null +++ b/examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Launch vLLM-Omni server for higgs-audio v3. +# +# Supports plain text TTS and voice cloning via /v1/audio/speech. +# +# Usage: +# ./run_server.sh # default port 8095, GPU 0 +# PORT=8096 GPUS=0,1 ./run_server.sh +# MODEL=/path/to/local/checkpoint ./run_server.sh + +set -e + +MODEL="${MODEL:-bosonai/higgs-audio-v3-tts-4b}" +PORT="${PORT:-8095}" +GPUS="${GPUS:-0}" +GPU_UTIL="${GPU_UTIL:-0.6}" + +echo "Starting higgs-audio v3 server" +echo " MODEL=$MODEL" +echo " PORT=$PORT" +echo " CUDA_VISIBLE_DEVICES=$GPUS" + +CUDA_VISIBLE_DEVICES="$GPUS" \ +VLLM_USE_DEEP_GEMM=0 \ +VLLM_MOE_USE_DEEP_GEMM=0 \ +vllm-omni serve "$MODEL" \ + --deploy-config vllm_omni/deploy/higgs_multimodal_qwen3.yaml \ + --host 0.0.0.0 \ + --port "$PORT" \ + --gpu-memory-utilization "$GPU_UTIL" \ + --trust-remote-code \ + --omni diff --git a/examples/online_serving/text_to_speech/ming_tts/README.md b/examples/online_serving/text_to_speech/ming_tts/README.md new file mode 100644 index 000000000..d0e0b0194 --- /dev/null +++ b/examples/online_serving/text_to_speech/ming_tts/README.md @@ -0,0 +1,113 @@ +# Ming-omni-tts Online Serving + +Serve the dense `inclusionAI/Ming-omni-tts-0.5B` two-stage TTS model through +the OpenAI-compatible `/v1/audio/speech` endpoint. + +## Start Server + +```bash +vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --omni \ + --port 8091 \ + --enforce-eager +``` + +Or: + +```bash +cd examples/online_serving/text_to_speech/ming_tts +./run_server.sh +``` + +The tested ROCm environment is summarized in the +[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md). + +## Send Requests + +The Python client targets `http://localhost:8091/v1` with `api_key=EMPTY`; it +does not call OpenAI's hosted API. + +```bash +python openai_speech_client.py \ + --text "你好，这是 Ming 在线语音合成测试。" \ + --max-new-tokens 200 +``` + +Style or dialect controls can be plain text or Ming JSON. The upstream +dialect example also uses `yue_prompt.wav` for speaker conditioning: + +```bash +python openai_speech_client.py \ + --text "我觉得社会企业同个人都有责任" \ + --instruction-json '{"方言":"广粤话"}' \ + --ref-audio /path/to/yue_prompt.wav \ + --max-new-tokens 200 +``` + +When `--ref-audio` is supplied without `--ref-text`, the server extracts the +Ming speaker embedding, matching upstream `use_spk_emb=True`, without using the +audio as a zero-shot prompt. + +Reference-audio cloning: + +```bash +python openai_speech_client.py \ + --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \ + --ref-audio /path/to/10002287-00000094.wav \ + --ref-text "在此奉劝大家别乱打美白针。" \ + --max-new-tokens 200 +``` + +Podcast-style multi-speaker prompt: + +```bash +python openai_speech_client.py \ + --text " speaker_1:你可以说一下，就大概说一下，可能虽然我也不知道，我看过那部电影没有。 + speaker_2:就是那个叫什么，变相一节课的嘛。 + speaker_1:嗯。 + speaker_2:一部搞笑的电影。 + speaker_1:一部搞笑的。" \ + --ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav \ + --ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \ + --ref-text " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试，做个，当服务员还要去笔试了 + speaker_2:对啊，这真的很奇怪，就是 单纯的因，单纯自己工资不高，只是因为可能人家那个店比较出名一点，就对你苛刻要求" +``` + +Streaming PCM: + +```bash +python openai_speech_client.py \ + --text "你好，这是流式输出测试。" \ + --stream \ + --output ming_output.pcm +``` + +`run_curl.sh` keeps small smoke checks: + +```bash +./run_curl.sh basic +REF_AUDIO=/path/to/reference.wav REF_TEXT="在此奉劝大家别乱打美白针。" ./run_curl.sh zero_shot +./run_curl.sh stream +``` + +## Request Fields + +| Field | Ming meaning | +|-------|--------------| +| `input` | target text | +| `instructions` | plain style text, or JSON object for structured Ming controls | +| `voice` | Ming IP voice label unless it resolves to an uploaded speaker | +| `language` | Ming `方言` control | +| `ref_audio` | speaker reference; with `ref_text`, also supplies the prompt waveform | +| `ref_text` | transcript enabling zero-shot or podcast prompt-latent conditioning | +| `speaker_embedding` | 192-d Ming speaker embedding | +| `max_new_tokens` | Ming `max_decode_steps` | + +## Notes + +- `ref_audio` accepts local paths through the client, remote URLs, `file://`, + or `data:` URLs. +- Non-streaming responses return WAV bytes; streaming responses return PCM. +- Music-only `bgm` generation is offline-only until the API exposes Ming + prompt-mode selection. diff --git a/examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py b/examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py new file mode 100644 index 000000000..536bfb048 --- /dev/null +++ b/examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech. + +Examples: + python openai_speech_client.py --text "你好，世界" + python openai_speech_client.py --text "我觉得社会企业同个人都有责任" \ + --instruction-json '{"方言":"广粤话"}' --ref-audio yue_prompt.wav --max-new-tokens 200 + python openai_speech_client.py \ + --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \ + --ref-audio 10002287-00000094.wav --ref-text "在此奉劝大家别乱打美白针。" \ + --max-new-tokens 200 + python openai_speech_client.py --text "speaker_1:你好。 speaker_2:你好。" \ + --ref-audio speaker_1.wav --ref-audio speaker_2.wav --ref-text "speaker_1:你好。 speaker_2:你好。" + python openai_speech_client.py --text "你好，这是流式输出测试。" \ + --stream --output ming_output.pcm +""" + +import argparse +import base64 +import json +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B" +EXPECTED_SPEAKER_EMBEDDING_DIM = 192 + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file to a base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_map = { + "wav": "audio/wav", + "mp3": "audio/mpeg", + "flac": "audio/flac", + "ogg": "audio/ogg", + "aac": "audio/aac", + } + mime_type = mime_map.get(ext, "audio/wav") + with open(audio_path, "rb") as f: + audio_b64 = base64.b64encode(f.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def load_speaker_embedding(path: str) -> list[float]: + """Load and validate a 192-d Ming speaker embedding JSON file.""" + with open(path, encoding="utf-8") as f: + data = json.load(f) + + if not isinstance(data, list): + raise ValueError("speaker_embedding file must contain a JSON list") + if len(data) != EXPECTED_SPEAKER_EMBEDDING_DIM: + raise ValueError( + f"Ming dense speaker_embedding must have {EXPECTED_SPEAKER_EMBEDDING_DIM} values, got {len(data)}" + ) + + values = [] + for index, value in enumerate(data): + try: + values.append(float(value)) + except (TypeError, ValueError) as exc: + raise ValueError(f"speaker_embedding[{index}] must be a number, got {value!r}") from exc + return values + + +def build_instruction_payload(args) -> str | None: + """Return a string payload for the API `instructions` field.""" + if args.instructions and args.instruction_json: + raise ValueError("Use either --instructions or --instruction-json, not both") + if args.instruction_json: + parsed = json.loads(args.instruction_json) + return json.dumps(parsed, ensure_ascii=False) + return args.instructions + + +def validate_args(args) -> None: + """Fail fast on invalid combinations before hitting the server.""" + if args.ref_text and not args.ref_audio: + raise ValueError("--ref-audio is required when --ref-text is provided") + if args.speaker_embedding and args.ref_audio and len(args.ref_audio) > 1: + raise ValueError("--speaker-embedding cannot be combined with multiple --ref-audio values") + + +def run_tts(args) -> None: + """Generate speech via the OpenAI-compatible /v1/audio/speech API.""" + validate_args(args) + + payload = { + "model": args.model, + "input": args.text, + "response_format": args.response_format, + } + + if args.voice: + payload["voice"] = args.voice + if args.task_type: + payload["task_type"] = args.task_type + if args.dialect: + payload["language"] = args.dialect + + instructions = build_instruction_payload(args) + if instructions: + payload["instructions"] = instructions + + if args.ref_audio: + ref_audio = [] + for audio in args.ref_audio: + if audio.startswith(("http://", "https://", "data:")): + ref_audio.append(audio) + else: + ref_audio.append(encode_audio_to_base64(audio)) + payload["ref_audio"] = ref_audio[0] if len(ref_audio) == 1 else ref_audio + if args.ref_text: + payload["ref_text"] = args.ref_text + if args.speaker_embedding: + payload["speaker_embedding"] = load_speaker_embedding(args.speaker_embedding) + if args.max_new_tokens: + payload["max_new_tokens"] = args.max_new_tokens + if args.stream: + payload["stream"] = True + payload["response_format"] = "pcm" + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + print(f"Payload keys: {sorted(payload)}") + + if args.stream: + output_path = args.output or "ming_output.pcm" + with httpx.Client(timeout=300.0) as client: + with client.stream("POST", api_url, json=payload, headers=headers) as response: + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.read().decode()) + return + with open(output_path, "wb") as f: + for chunk in response.iter_bytes(): + f.write(chunk) + print(f"Streamed PCM audio to: {output_path}") + return + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "ming_output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument( + "--task-type", + default=None, + choices=["CustomVoice", "VoiceDesign", "Base"], + help="Optional compatibility task type. Ming accepts the same field but primarily uses prompt metadata.", + ) + parser.add_argument( + "--voice", + default=None, + help="Maps to Ming `IP` when using built-in character voices, or to an uploaded voice sample name", + ) + parser.add_argument("--dialect", default=None, help="Maps to Ming `方言`") + parser.add_argument("--instructions", default=None, help="Free-form Ming instruction string") + parser.add_argument( + "--instruction-json", + default=None, + help='Structured Ming instruction JSON, for example \'{"情感":"高兴"}\'', + ) + parser.add_argument( + "--ref-audio", + action="append", + default=None, + help="Reference audio path, URL, or data URL. Repeat for podcast-style multi-speaker prompts.", + ) + parser.add_argument("--ref-text", default=None, help="Reference transcript for cloning") + parser.add_argument( + "--speaker-embedding", default=None, help="Path to a JSON file containing a 192-d speaker embedding" + ) + parser.add_argument("--max-new-tokens", type=int, default=None, help="Override ming_max_decode_steps") + parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format when not streaming", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + try: + run_tts(args) + except (FileNotFoundError, ValueError, json.JSONDecodeError) as exc: + raise SystemExit(f"Error: {exc}") from exc + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_speech/ming_tts/run_curl.sh b/examples/online_serving/text_to_speech/ming_tts/run_curl.sh new file mode 100755 index 000000000..10fcaf59a --- /dev/null +++ b/examples/online_serving/text_to_speech/ming_tts/run_curl.sh @@ -0,0 +1,76 @@ +#!/bin/bash +set -euo pipefail + +MODE="${1:-basic}" +HOST="${HOST:-localhost}" +PORT="${PORT:-8091}" +MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}" +API_URL="http://${HOST}:${PORT}/v1/audio/speech" +TEXT="${TEXT:-你好，这是 Ming 在线语音合成测试。}" +OUTPUT="${OUTPUT:-ming_output.wav}" +STREAM_OUTPUT="${STREAM_OUTPUT:-ming_output.pcm}" +REF_AUDIO="${REF_AUDIO:-}" +REF_TEXT="${REF_TEXT:-}" + +post_json() { + local payload="$1" + local output_path="$2" + curl -X POST "$API_URL" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer EMPTY" \ + -d "$payload" \ + --output "$output_path" +} + +case "$MODE" in + basic) + post_json "{ + \"model\": \"${MODEL}\", + \"input\": \"${TEXT}\", + \"response_format\": \"wav\" + }" "$OUTPUT" + ;; + zero_shot) + if [ -z "$REF_AUDIO" ] || [ -z "$REF_TEXT" ]; then + echo "zero_shot requires REF_AUDIO and REF_TEXT" >&2 + exit 1 + fi + python - <<'PY' > /tmp/ming_zero_shot_payload.json +import base64 +import json +import mimetypes +import os +from pathlib import Path + +path = Path(os.environ["REF_AUDIO"]) +mime_type = mimetypes.guess_type(path.name)[0] or "audio/wav" +payload = { + "model": os.environ["MODEL"], + "input": os.environ["TEXT"], + "ref_audio": f"data:{mime_type};base64,{base64.b64encode(path.read_bytes()).decode('utf-8')}", + "ref_text": os.environ["REF_TEXT"], + "response_format": "wav", +} +print(json.dumps(payload, ensure_ascii=False)) +PY + curl -X POST "$API_URL" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer EMPTY" \ + --data-binary @/tmp/ming_zero_shot_payload.json \ + --output "$OUTPUT" + rm -f /tmp/ming_zero_shot_payload.json + ;; + stream) + post_json "{ + \"model\": \"${MODEL}\", + \"input\": \"${TEXT}\", + \"stream\": true, + \"response_format\": \"pcm\" + }" "$STREAM_OUTPUT" + ;; + *) + echo "Unknown mode: $MODE" >&2 + echo "Supported sanity checks: basic, zero_shot, stream" >&2 + exit 1 + ;; +esac diff --git a/examples/online_serving/text_to_speech/ming_tts/run_server.sh b/examples/online_serving/text_to_speech/ming_tts/run_server.sh new file mode 100755 index 000000000..7d294d950 --- /dev/null +++ b/examples/online_serving/text_to_speech/ming_tts/run_server.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Launch vLLM-Omni server for Ming-omni-tts. +# +# Usage: +# ./run_server.sh +# PORT=8000 ./run_server.sh + +set -e + +DIR="$(cd "$(dirname "$0")" && pwd)" +ROOT="$(cd "$DIR/../../../.." && pwd)" + +MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}" +PORT="${PORT:-8091}" +DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT/vllm_omni/deploy/ming_tts.yaml}" + +echo "Starting Ming-omni-tts server with model: $MODEL" +echo "Deploy config: $DEPLOY_CONFIG" + +vllm-omni serve "$MODEL" \ + --deploy-config "$DEPLOY_CONFIG" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --enforce-eager \ + --omni diff --git a/examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py b/examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py index 8b39b9eab..3c99a82d3 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py @@ -408,7 +408,10 @@ def create_app(api_base: str): await resp.aclose() await client.aclose() - return StreamingResponse(relay(), media_type="application/octet-stream") + # Use audio/pcm to bypass browsers (e.g. Chromium) MIME type sniffing, which + # buffers application/octet-stream responses before delivering chunks to + # fetch().body.getReader(), inflating TTFP to total_ms. + return StreamingResponse(relay(), media_type="audio/pcm") # ── Gradio UI ──────────────────────────────────────────────── voices = fetch_voices(api_base) diff --git a/examples/quantization/check_modelopt_fp8_export.py b/examples/quantization/check_modelopt_fp8_export.py new file mode 100644 index 000000000..90d952a8f --- /dev/null +++ b/examples/quantization/check_modelopt_fp8_export.py @@ -0,0 +1,247 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Verify a ModelOpt FP8 diffusers checkpoint exported by +quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py. + +Three checks: + A. transformer/config.json has a sane quantization_config block. + B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors. + C. transformer disk size is materially smaller than a BF16 baseline. + +Example: + python examples/quantization/check_modelopt_fp8_export.py \\ + --output ./hv15-480p-modelopt-fp8 + + # Optional: compare disk size against a local or HF BF16 baseline. + python examples/quantization/check_modelopt_fp8_export.py \\ + --output ./hv15-480p-modelopt-fp8 \\ + --baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v +""" + +from __future__ import annotations + +import argparse +import json +import sys +from collections import Counter +from pathlib import Path + +SUPPORTED_ALGOS = ("FP8",) + + +def _check_config(transformer_dir: Path) -> tuple[int, str | None]: + """Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn.""" + cfg_path = transformer_dir / "config.json" + if not cfg_path.exists(): + print(f"[FAIL] {cfg_path} missing.") + return 1, None + + with cfg_path.open(encoding="utf-8") as f: + cfg = json.load(f) + + qc = cfg.get("quantization_config") + if not isinstance(qc, dict): + print(f"[FAIL] No `quantization_config` block in {cfg_path}.") + return 1, None + + print(f"[A] quantization_config from {cfg_path}:") + print(json.dumps(qc, indent=2)) + + quant_algo = qc.get("quant_algo") + issues = [] + if qc.get("quant_method") != "modelopt": + issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')") + if quant_algo not in SUPPORTED_ALGOS: + issues.append( + f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)" + ) + + if issues: + print("[A] WARN — config looks incomplete:") + for issue in issues: + print(f" - {issue}") + return 2, quant_algo + print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).") + return 0, quant_algo + + +def _read_safetensors_header(path: Path) -> dict: + """Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors. + + Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}. + Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ... + """ + import struct + + with open(path, "rb") as f: + header_len = struct.unpack(" str: + """Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes. + + ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1 + interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 — + and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block. + """ + if not weight_scale_shapes: + return "no weight_scale tensors found" + + def meaningful_dims(shape: list[int]) -> int: + return sum(1 for d in shape if d > 1) + + per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0) + per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1) + per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2) + total = len(weight_scale_shapes) + if per_tensor == total: + return "per-tensor (all scalar scales)" + if per_channel == total: + return "per-channel (1 meaningful dim)" + if per_block == total: + return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)" + return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}" + + +def _check_safetensors(transformer_dir: Path) -> int: + """Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header.""" + files = sorted(transformer_dir.glob("*.safetensors")) + if not files: + print(f"[FAIL] No *.safetensors in {transformer_dir}.") + return 1 + + header_dtype_counts: Counter[str] = Counter() + sample_quant_weight_keys: list[str] = [] + sample_scale_keys: list[str] = [] + weight_scale_shapes: list[list[int]] = [] + sample_weight_scale_entries: list[tuple[str, list[int]]] = [] + + for f in files: + try: + header = _read_safetensors_header(f) + except Exception as exc: + print(f"[B] WARN — could not parse header of {f}: {exc}") + continue + for k, info in header.items(): + dtype = info.get("dtype", "?") + header_dtype_counts[dtype] += 1 + # FP8 stores weights as F8_E4M3 directly. + if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5: + sample_quant_weight_keys.append(k) + if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5: + sample_scale_keys.append(k) + if k.endswith(".weight_scale"): + weight_scale_shapes.append(info.get("shape", [])) + if len(sample_weight_scale_entries) < 5: + sample_weight_scale_entries.append((k, info.get("shape", []))) + + print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):") + for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]): + marker = " <-- FP8" if dtype.startswith("F8") else "" + print(f" {dtype:10s} {count:>6d}{marker}") + + quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8")) + if quant_count == 0: + print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.") + return 1 + + print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.") + if sample_quant_weight_keys: + print(f" sample quantized weight tensors: {sample_quant_weight_keys[:3]}") + if sample_scale_keys: + print(f" sample scale tensors: {sample_scale_keys[:3]}") + print(" (Note: torch's get_tensor() may return these as bf16 views on some versions —") + print(" irrelevant; vLLM's loader uses native FP8 ops.)") + + # Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D). + print(f"\n weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}") + for key, shape in sample_weight_scale_entries[:3]: + print(f" {key}: shape {shape}") + return 0 + + +def _disk_size_gib(p: Path) -> float: + return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3) + + +def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int: + """Returns 0 always (informational only).""" + quant_size = _disk_size_gib(transformer_dir) + print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB") + + if baseline is None: + print("[C] SKIP — pass --baseline to compare against BF16.") + return 0 + + baseline_path = Path(baseline) + if not baseline_path.exists(): + # Try HF download. + try: + from huggingface_hub import snapshot_download + except ImportError: + print("[C] SKIP — huggingface_hub not installed and baseline not a local path.") + return 0 + print(f" Downloading baseline transformer from HF: {baseline}") + baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"])) + + bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path + bf16_size = _disk_size_gib(bf16_dir) + if bf16_size == 0: + print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}") + return 0 + + # Expected reduction for FP8: ~50%. + min_reduction = 30 + reduction = (1 - quant_size / bf16_size) * 100 + print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})") + print(f"[C] Disk reduction: {reduction:.1f}% (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)") + if reduction < min_reduction: + print( + f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization." + ) + return 0 + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.") + p.add_argument( + "--baseline", + default=None, + help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.", + ) + args = p.parse_args() + + out_root = Path(args.output).expanduser().resolve() + transformer_dir = out_root / "transformer" + if not transformer_dir.exists(): + print(f"[FAIL] {transformer_dir} does not exist.") + sys.exit(1) + + print(f"Checking: {out_root}\n") + + fail = 0 + config_status, _quant_algo = _check_config(transformer_dir) + fail |= config_status + fail |= _check_safetensors(transformer_dir) + _check_size_vs_baseline(transformer_dir, args.baseline) + + print() + if fail == 0: + print("=" * 60) + print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.") + elif fail == 1: + print("=" * 60) + print("FAILURES detected — calibration may need to be re-run.") + sys.exit(1) + else: + print("=" * 60) + print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.") + + +if __name__ == "__main__": + main() diff --git a/examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py b/examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py new file mode 100644 index 000000000..e6e3dd5c5 --- /dev/null +++ b/examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint. + +Calibrates the DiT transformer using a small video prompt set and exports a +diffusers-style directory whose transformer carries ModelOpt FP8 metadata. +The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter +(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py). + +Layers kept full precision match the #2728 / #2795 pattern: modulation, +AdaLayerNorm, entry/exit projections, embeddings, the token refiner path, +and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention +empirically degrades under FP8 (see #2920 ablation). + +Example: + python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\ + --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\ + --output ./hv15-480p-modelopt-fp8 \\ + --overwrite +""" + +from __future__ import annotations + +import argparse +import copy +import json +import re +import shutil +import sys +from pathlib import Path +from typing import Any + +import torch +from diffusers import DiffusionPipeline + +DEFAULT_PROMPTS = [ + "A dog running across a field of golden wheat.", + "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.", + "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.", + "A crackling campfire at night under a starry sky, sparks rising into the dark.", + "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.", + "A close-up of a blooming rose covered in morning dew, soft natural light.", + "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.", + "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.", +] + + +def _build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.") + p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.") + p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16") + p.add_argument("--height", type=int, default=480) + p.add_argument("--width", type=int, default=832) + p.add_argument( + "--num-frames", + type=int, + default=33, + help="Frames per calibration sample. 33 matches the typical short benchmark.", + ) + p.add_argument("--guidance-scale", type=float, default=6.0) + p.add_argument( + "--calib-steps", + type=int, + default=10, + help="Denoising steps per calibration prompt (10 is enough for amax statistics).", + ) + p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.") + p.add_argument("--seed", type=int, default=42) + p.add_argument( + "--prompt", + action="append", + default=[], + help="Custom calibration prompt. Repeat to provide multiple.", + ) + p.add_argument( + "--quantize-mha", + action="store_true", + help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.", + ) + p.add_argument( + "--weight-block-size", + type=str, + default=None, + help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). " + "Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at " + "negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.", + ) + p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.") + return p + + +def _parse_block_size(spec: str | None) -> list[int] | None: + if spec is None: + return None + parts = [int(x) for x in spec.split(",") if x.strip()] + if len(parts) != 2: + raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}") + return parts + + +def _require_modelopt() -> Any: + try: + import modelopt.torch.quantization as mtq + except ModuleNotFoundError as exc: + raise SystemExit( + "NVIDIA ModelOpt is not installed. Install with:\n" + " pip install 'nvidia-modelopt[all]'\n" + f"Original error: {exc}" + ) from exc + return mtq + + +def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]: + model_path = args.model + output_dir = Path(args.output).expanduser().resolve() + if output_dir.exists(): + if not args.overwrite: + raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.") + shutil.rmtree(output_dir) + return model_path, output_dir + + +def _select_dtype(name: str) -> torch.dtype: + return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name] + + +def _build_prompts(args: argparse.Namespace) -> list[str]: + prompts = args.prompt or DEFAULT_PROMPTS + if args.calib_size <= 0: + raise SystemExit("--calib-size must be positive.") + if len(prompts) < args.calib_size: + repeats = (args.calib_size + len(prompts) - 1) // len(prompts) + prompts = (prompts * repeats)[: args.calib_size] + return prompts[: args.calib_size] + + +# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern). +# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding +# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation +# - proj_out: final output projection +# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear +def _filter_func_hv15(name: str) -> bool: + pattern = re.compile( + r"(proj_out.*|" + r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|" + r"time_embed|cond_type_embed|" + r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|" + r"token_refiner).*)" + ) + return pattern.match(name) is not None + + +def _mha_filter_func(name: str) -> bool: + pattern = re.compile( + r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*" + ) + return pattern.match(name) is not None + + +def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None: + if not hasattr(mtq, "disable_quantizer"): + return + mtq.disable_quantizer(backbone, _filter_func_hv15) + if not quantize_mha: + mtq.disable_quantizer(backbone, _mha_filter_func) + + +def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline: + pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype) + if hasattr(pipe, "set_progress_bar_config"): + pipe.set_progress_bar_config(disable=True) + pipe.to("cuda") + return pipe + + +def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]): + generator = torch.Generator(device="cuda") + + # Try to set guidance on the pipeline's guider object up front (modern + # diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls + # back silently — calibration uses whatever default the pipeline ships with. + guider = getattr(pipe, "guider", None) + if guider is not None and hasattr(guider, "guidance_scale"): + try: + guider.guidance_scale = args.guidance_scale + except Exception: + pass + + base_kwargs = dict( + height=args.height, + width=args.width, + num_frames=args.num_frames, + num_inference_steps=args.calib_steps, + output_type="latent", + ) + + def forward_loop(*_unused_args, **_unused_kwargs) -> None: + with torch.inference_mode(): + for idx, prompt in enumerate(prompts): + generator.manual_seed(args.seed + idx) + # Try with guidance_scale first; fall back without on TypeError + # for pipelines (like HV-1.5) that take CFG via guider config. + try: + pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs) + except TypeError as exc: + if "guidance_scale" not in str(exc): + raise + pipe(prompt=prompt, generator=generator, **base_kwargs) + + return forward_loop + + +def _summarize_export(output_dir: Path) -> None: + cfg_path = output_dir / "transformer" / "config.json" + if not cfg_path.exists(): + print(f"[warn] {cfg_path} missing.", file=sys.stderr) + return + with cfg_path.open(encoding="utf-8") as f: + cfg = json.load(f) + qc = cfg.get("quantization_config") + if not isinstance(qc, dict): + print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr) + return + print("Export summary:") + print(f" quant_method: {qc.get('quant_method')}") + print(f" quant_algo: {qc.get('quant_algo')}") + producer = qc.get("producer") + if isinstance(producer, dict): + print(f" producer: {producer.get('name')} {producer.get('version')}") + print(f" config path: {cfg_path}") + + +def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int: + """Convert in-memory weights of quantized modules to actual FP8 storage. + + `export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't + in ModelOpt's recognized-model registry), so we must call the per-weight + export helper ourselves. Same workaround as the HunyuanImage-3 calibration + helper. + """ + from modelopt.torch.export.quant_utils import ( + QUANTIZATION_NONE, + get_quantization_format, + quantizer_attr_names, + weight_attr_names, + ) + from modelopt.torch.export.unified_export_hf import _export_quantized_weight + + exported = 0 + for name, module in backbone.named_modules(): + try: + quantization_format = get_quantization_format(module) + except Exception as exc: + print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr) + continue + if quantization_format == QUANTIZATION_NONE: + continue + for weight_name in weight_attr_names(module): + quantizer_attrs = quantizer_attr_names(weight_name) + weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None) + if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False): + continue + _export_quantized_weight(module, dtype, weight_name) + exported += 1 + return exported + + +def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict: + """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913). + + Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When + `weight_block_size` is given, advertise block-wise weight quantization in + the saved metadata (so consumers know to expect multi-element scale tensors). + """ + weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"} + if weight_block_size is not None: + weights_cfg["strategy"] = "block" + weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}" + return { + "config_groups": { + "group_0": { + "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, + "weights": weights_cfg, + "targets": ["Linear"], + } + }, + "ignore": [ + "context_embedder*", + "context_embedder_2*", + "cond_type_embed*", + "image_embedder*", + "norm1.linear*", + "norm1_context.linear*", + "norm2*", + "norm2_context*", + "norm_out*", + "proj_out*", + "time_embed*", + "token_refiner*", + "x_embedder*", + ], + "producer": {"name": "modelopt"}, + "quant_algo": "FP8", + "quant_method": "modelopt", + } + + +def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None: + """Inject quant_algo: FP8 + config_groups into transformer/config.json so + vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.""" + cfg_path = output_dir / "transformer" / "config.json" + with cfg_path.open(encoding="utf-8") as f: + cfg = json.load(f) + + new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size) + existing = cfg.get("quantization_config") + if isinstance(existing, dict): + producer = existing.get("producer") + if isinstance(producer, dict): + new_qc["producer"] = producer + + cfg["quantization_config"] = new_qc + with cfg_path.open("w", encoding="utf-8") as f: + json.dump(cfg, f, indent=2) + + +def _save_pipeline_with_fp8_transformer( + pipe: DiffusionPipeline, + model_path: str, + output_dir: Path, + max_shard_size: str = "5GB", +) -> None: + """Save the pipeline with the (now FP8) transformer. + + Copies the source directory verbatim except for `transformer/`, then + saves the transformer with quantizers hidden so the state dict contains + only the FP8 weights + scale tensors. + """ + from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict + + src = Path(model_path) + if not src.exists(): + from huggingface_hub import snapshot_download + + src = Path(snapshot_download(model_path)) + + if output_dir.exists(): + shutil.rmtree(output_dir) + shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer")) + + transformer_out = output_dir / "transformer" + # `hide_quantizers_from_state_dict` walks named_modules(); pass the actual + # nn.Module (transformer), not the diffusers Pipeline wrapper. + with hide_quantizers_from_state_dict(pipe.transformer): + pipe.transformer.save_pretrained( + str(transformer_out), + safe_serialization=True, + max_shard_size=max_shard_size, + ) + + +def main() -> None: + args = _build_parser().parse_args() + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required for ModelOpt FP8 quantization.") + + mtq = _require_modelopt() + model_path, output_dir = _ensure_paths(args) + dtype = _select_dtype(args.dtype) + prompts = _build_prompts(args) + + print("Quantization plan:") + weight_block_size = _parse_block_size(args.weight_block_size) + + print(f" input: {args.model}") + print(f" output: {output_dir}") + print(f" dtype: {dtype}") + print(f" height/width: {args.height}x{args.width}") + print(f" num_frames: {args.num_frames}") + print(f" calib_size: {len(prompts)}") + print(f" calib_steps: {args.calib_steps}") + print(f" quantize_mha: {args.quantize_mha}") + print( + f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}" + ) + + pipe = _load_pipeline(model_path, dtype) + backbone = pipe.transformer + + quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + if weight_block_size is not None: + # Switch from per-tensor (default) to block-wise weight quantization. + # ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer. + quant_config["quant_cfg"]["*weight_quantizer"] = { + "num_bits": (4, 3), # E4M3 (FP8 weights, same as default) + "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]}, + } + print( + f" -> overriding weight quantizer with block_sizes={weight_block_size} " + f"({weight_block_size[0]}x{weight_block_size[1]} tiles)" + ) + + forward_loop = _build_forward_loop(pipe, args, prompts) + quantized = mtq.quantize(backbone, quant_config, forward_loop) + if quantized is not None: + pipe.transformer = quantized + backbone = quantized + + _disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha) + + print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,") + print("so we have to call the per-weight export helper ourselves)...") + exported = _force_export_quantized_weights(backbone, dtype) + print(f" -> {exported} weights converted to FP8 in memory") + if exported == 0: + raise SystemExit( + "No quantized weights were exported. Calibration may have skipped every layer " + "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any " + "weight quantizers." + ) + + print("\nSaving pipeline with FP8 transformer...") + _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir) + _patch_quant_config(output_dir, weight_block_size=weight_block_size) + print(f"Saved to: {output_dir}") + _summarize_export(output_dir) + + print("\nNext: validate the checkpoint with vllm-omni:") + print( + " python examples/offline_inference/text_to_video/text_to_video.py \\\n" + f" --model {output_dir} \\\n" + " --quantization fp8 \\\n" + " --prompt 'A dog running across a field of golden wheat.' \\\n" + f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n" + " --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n" + " --output outputs/hv15_modelopt_fp8.mp4 \\\n" + " --enforce-eager" + ) + print( + "\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the " + "checkpoint's config.json has modelopt metadata.)" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/quantization/quantize_wan2_2_modelopt_fp8.py b/examples/quantization/quantize_wan2_2_modelopt_fp8.py new file mode 100644 index 000000000..6b6cc42a4 --- /dev/null +++ b/examples/quantization/quantize_wan2_2_modelopt_fp8.py @@ -0,0 +1,471 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint. + +Calibrates the DiT transformer using a small video prompt set and exports a +diffusers-style directory whose transformer carries ModelOpt FP8 metadata. +The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter +(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py). + +Layers kept full precision match the #2728 / #2795 pattern: condition embedder +(time/text/image), patch embedding, modulation (scale_shift_table), final +norm + proj_out, and sequence-parallel helpers. All attention + FFN linears +are quantized — static calibration handles the numerics that online FP8 +couldn't (see #2920 ablation). + +Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that +fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here. + +Example: + python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\ + --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\ + --output ./wan22-ti2v-modelopt-fp8 \\ + --overwrite +""" + +from __future__ import annotations + +import argparse +import copy +import json +import re +import shutil +import sys +from pathlib import Path +from typing import Any + +import torch +from diffusers import DiffusionPipeline + +DEFAULT_PROMPTS = [ + "A dog running across a field of golden wheat.", + "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.", + "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.", + "A crackling campfire at night under a starry sky, sparks rising into the dark.", + "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.", + "A close-up of a blooming rose covered in morning dew, soft natural light.", + "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.", + "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.", +] + + +def _build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.") + p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.") + p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16") + p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).") + p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).") + p.add_argument( + "--num-frames", + type=int, + default=49, + help="Frames per calibration sample. 49 matches the typical short benchmark; " + "use 17 to reduce memory pressure during calibration.", + ) + p.add_argument("--guidance-scale", type=float, default=5.0) + p.add_argument( + "--calib-steps", + type=int, + default=10, + help="Denoising steps per calibration prompt (10 is enough for amax statistics).", + ) + p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.") + p.add_argument("--seed", type=int, default=42) + p.add_argument( + "--prompt", + action="append", + default=[], + help="Custom calibration prompt. Repeat to provide multiple.", + ) + p.add_argument( + "--quantize-mha", + action="store_true", + help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention " + "sequences amplified FP8 drift in the online ablation (see #2920).", + ) + p.add_argument( + "--weight-block-size", + type=str, + default=None, + help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. " + "Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 " + "for the HV-1.5 investigation status before relying on this.", + ) + p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.") + return p + + +def _parse_block_size(spec: str | None) -> list[int] | None: + if spec is None: + return None + parts = [int(x) for x in spec.split(",") if x.strip()] + if len(parts) != 2: + raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}") + return parts + + +def _require_modelopt() -> Any: + try: + import modelopt.torch.quantization as mtq + except ModuleNotFoundError as exc: + raise SystemExit( + "NVIDIA ModelOpt is not installed. Install with:\n" + " pip install 'nvidia-modelopt[all]'\n" + f"Original error: {exc}" + ) from exc + return mtq + + +def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]: + model_path = args.model + output_dir = Path(args.output).expanduser().resolve() + if output_dir.exists(): + if not args.overwrite: + raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.") + shutil.rmtree(output_dir) + return model_path, output_dir + + +def _select_dtype(name: str) -> torch.dtype: + return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name] + + +def _build_prompts(args: argparse.Namespace) -> list[str]: + prompts = args.prompt or DEFAULT_PROMPTS + if args.calib_size <= 0: + raise SystemExit("--calib-size must be positive.") + if len(prompts) < args.calib_size: + repeats = (args.calib_size + len(prompts) - 1) // len(prompts) + prompts = (prompts * repeats)[: args.calib_size] + return prompts[: args.calib_size] + + +# Layers to KEEP at full precision. Wan2.2's module naming: +# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V) +# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip) +# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard) +# - norm_out: AdaLayerNorm final +# - proj_out: final nn.Linear +# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers +def _filter_func_wan22(name: str) -> bool: + pattern = re.compile( + r"(proj_out.*|" + r".*(condition_embedder|patch_embedding|" + r"norm_out|scale_shift_table|" + r"timestep_proj_prepare|output_scale_shift_prepare).*)" + ) + return pattern.match(name) is not None + + +def _mha_filter_func(name: str) -> bool: + pattern = re.compile( + r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*" + ) + return pattern.match(name) is not None + + +def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None: + if not hasattr(mtq, "disable_quantizer"): + return + mtq.disable_quantizer(backbone, _filter_func_wan22) + if not quantize_mha: + mtq.disable_quantizer(backbone, _mha_filter_func) + + +def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline: + pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype) + if hasattr(pipe, "set_progress_bar_config"): + pipe.set_progress_bar_config(disable=True) + pipe.to("cuda") + return pipe + + +def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]): + generator = torch.Generator(device="cuda") + + # Try setting guidance on the pipeline's guider if present (newer diffusers APIs). + guider = getattr(pipe, "guider", None) + if guider is not None and hasattr(guider, "guidance_scale"): + try: + guider.guidance_scale = args.guidance_scale + except Exception: + pass + + base_kwargs = dict( + height=args.height, + width=args.width, + num_frames=args.num_frames, + num_inference_steps=args.calib_steps, + output_type="latent", + ) + + def forward_loop(*_unused_args, **_unused_kwargs) -> None: + with torch.inference_mode(): + for idx, prompt in enumerate(prompts): + generator.manual_seed(args.seed + idx) + # Try with guidance_scale first; fall back without on TypeError + # for pipelines that take CFG via guider config only. + try: + pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs) + except TypeError as exc: + if "guidance_scale" not in str(exc): + raise + pipe(prompt=prompt, generator=generator, **base_kwargs) + + return forward_loop + + +def _summarize_export(output_dir: Path) -> None: + cfg_path = output_dir / "transformer" / "config.json" + if not cfg_path.exists(): + print(f"[warn] {cfg_path} missing.", file=sys.stderr) + return + with cfg_path.open(encoding="utf-8") as f: + cfg = json.load(f) + qc = cfg.get("quantization_config") + if not isinstance(qc, dict): + print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr) + return + print("Export summary:") + print(f" quant_method: {qc.get('quant_method')}") + print(f" quant_algo: {qc.get('quant_algo')}") + producer = qc.get("producer") + if isinstance(producer, dict): + print(f" producer: {producer.get('name')} {producer.get('version')}") + print(f" config path: {cfg_path}") + + +def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int: + """Convert in-memory weights of quantized modules to actual FP8 storage. + + `export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't + in ModelOpt's recognized-model registry), so we must call the per-weight + export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3 + calibration helpers. + """ + from modelopt.torch.export.quant_utils import ( + QUANTIZATION_NONE, + get_quantization_format, + quantizer_attr_names, + weight_attr_names, + ) + from modelopt.torch.export.unified_export_hf import _export_quantized_weight + + exported = 0 + for name, module in backbone.named_modules(): + try: + quantization_format = get_quantization_format(module) + except Exception as exc: + print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr) + continue + if quantization_format == QUANTIZATION_NONE: + continue + for weight_name in weight_attr_names(module): + quantizer_attrs = quantizer_attr_names(weight_name) + weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None) + if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False): + continue + _export_quantized_weight(module, dtype, weight_name) + exported += 1 + return exported + + +def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict: + """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).""" + weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"} + if weight_block_size is not None: + weights_cfg["strategy"] = "block" + weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}" + return { + "config_groups": { + "group_0": { + "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, + "weights": weights_cfg, + "targets": ["Linear"], + } + }, + "ignore": [ + "condition_embedder*", + "norm_out*", + "output_scale_shift_prepare*", + "patch_embedding*", + "proj_out*", + "scale_shift_table*", + "timestep_proj_prepare*", + ], + "producer": {"name": "modelopt"}, + "quant_algo": "FP8", + "quant_method": "modelopt", + } + + +def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None: + """Inject quant_algo: FP8 + config_groups into each transformer's config.json + so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8. + + Patches both ``transformer`` and ``transformer_2`` (A14B) when present.""" + for sub in ("transformer", "transformer_2"): + cfg_path = output_dir / sub / "config.json" + if not cfg_path.exists(): + continue + with cfg_path.open(encoding="utf-8") as f: + cfg = json.load(f) + + new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size) + existing = cfg.get("quantization_config") + if isinstance(existing, dict): + producer = existing.get("producer") + if isinstance(producer, dict): + new_qc["producer"] = producer + + cfg["quantization_config"] = new_qc + with cfg_path.open("w", encoding="utf-8") as f: + json.dump(cfg, f, indent=2) + + +def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]: + """Return [(attr_name, module)] for every DiT on the pipeline. + + Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2`` + (high-noise). TI2V-5B has only ``transformer``. + """ + out = [] + for attr in ("transformer", "transformer_2"): + mod = getattr(pipe, attr, None) + if mod is not None: + out.append((attr, mod)) + return out + + +def _save_pipeline_with_fp8_transformer( + pipe: DiffusionPipeline, + model_path: str, + output_dir: Path, + max_shard_size: str = "5GB", +) -> None: + """Copy source dir verbatim minus transformer dirs, then save each quantized DiT. + + Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well + as the single-DiT TI2V-5B. + """ + from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict + + src = Path(model_path) + if not src.exists(): + from huggingface_hub import snapshot_download + + src = Path(snapshot_download(model_path)) + + if output_dir.exists(): + shutil.rmtree(output_dir) + shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2")) + + for attr, backbone in _list_transformers(pipe): + transformer_out = output_dir / attr + # Pass the nn.Module (transformer), not the Pipeline wrapper. + with hide_quantizers_from_state_dict(backbone): + backbone.save_pretrained( + str(transformer_out), + safe_serialization=True, + max_shard_size=max_shard_size, + ) + + +def main() -> None: + args = _build_parser().parse_args() + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required for ModelOpt FP8 quantization.") + + mtq = _require_modelopt() + model_path, output_dir = _ensure_paths(args) + dtype = _select_dtype(args.dtype) + prompts = _build_prompts(args) + weight_block_size = _parse_block_size(args.weight_block_size) + + print("Quantization plan:") + print(f" input: {args.model}") + print(f" output: {output_dir}") + print(f" dtype: {dtype}") + print(f" height/width: {args.height}x{args.width}") + print(f" num_frames: {args.num_frames}") + print(f" calib_size: {len(prompts)}") + print(f" calib_steps: {args.calib_steps}") + print(f" quantize_mha: {args.quantize_mha}") + print( + f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}" + ) + + pipe = _load_pipeline(model_path, dtype) + transformers = _list_transformers(pipe) + if not transformers: + raise SystemExit("Pipeline has no transformer or transformer_2 attribute.") + print(f" found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}") + + quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + if weight_block_size is not None: + quant_config["quant_cfg"]["*weight_quantizer"] = { + "num_bits": (4, 3), + "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]}, + } + print( + f" -> overriding weight quantizer with block_sizes={weight_block_size} " + f"({weight_block_size[0]}x{weight_block_size[1]} tiles)" + ) + + forward_loop = _build_forward_loop(pipe, args, prompts) + + # Quantize each DiT. The forward_loop runs the full pipeline, so quantizers + # on every attached transformer (A14B: transformer + transformer_2) get + # calibrated during the same passes. + for attr, backbone in transformers: + print(f"\nQuantizing {attr}...") + quantized = mtq.quantize(backbone, quant_config, forward_loop) + if quantized is not None: + setattr(pipe, attr, quantized) + _disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha) + + # Export weights to FP8 only AFTER all transformers are calibrated. Exporting + # inside the loop would convert transformer's weights to real FP8, and the + # next transformer's forward_loop (full pipeline) would then fake-quantize an + # already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn". + print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,") + print("so we have to call the per-weight export helper ourselves)...") + total_exported = 0 + for attr, backbone in _list_transformers(pipe): + exported = _force_export_quantized_weights(backbone, dtype) + print(f" -> {exported} weights converted to FP8 in {attr}") + total_exported += exported + + if total_exported == 0: + raise SystemExit( + "No quantized weights were exported. Calibration may have skipped every layer " + "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any " + "weight quantizers." + ) + + print("\nSaving pipeline with FP8 transformer(s)...") + _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir) + _patch_quant_config(output_dir, weight_block_size=weight_block_size) + print(f"Saved to: {output_dir}") + _summarize_export(output_dir) + + print("\nNext: validate the checkpoint with vllm-omni:") + print( + " python examples/offline_inference/text_to_video/text_to_video.py \\\n" + f" --model {output_dir} \\\n" + " --quantization fp8 \\\n" + " --prompt 'A dog running across a field of golden wheat.' \\\n" + f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n" + " --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n" + " --output outputs/wan22_modelopt_fp8.mp4" + ) + print( + "\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the " + "checkpoint's config.json has modelopt metadata.)" + ) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 0950b2682..1bb8fb2f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,7 +65,9 @@ dev = [ "orjson>=3.11.7", # to make v1/embeddings API fast "voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests "torchdiffeq>=0.2.5", # for covo_audio Test - "FlagEmbedding" # for hunyuan_image3 tests + "FlagEmbedding", # for hunyuan_image3 tests + "kernels<0.15.0", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0) + "s3tokenizer>=0.3.0" # for tts test ] demo = [ @@ -89,6 +91,12 @@ minicpmo = [ "stepaudio2-minicpmo", ] +# Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported +# as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md. +quack = [ + "quack-kernels>=0.3.11", +] + docs = [ "mkdocs>=1.5.0", "mkdocs-api-autonav", @@ -285,5 +293,9 @@ fullset = "fullset" Vai = "Vai" tockens = "tockens" CANN = "CANN" +# Step-Audio reasoning parser streams `` tokens split across chunks +# (e.g. ""). The fragment "thi" appears literally in its +# parser logic and tests and is not a typo of "the"/"this". +thi = "thi" ND = "ND" byt5 = "byt5" diff --git a/recipes/BosonAI/Higgs-Audio-V3-TTS.md b/recipes/BosonAI/Higgs-Audio-V3-TTS.md new file mode 100644 index 000000000..b924ac472 --- /dev/null +++ b/recipes/BosonAI/Higgs-Audio-V3-TTS.md @@ -0,0 +1,131 @@ +# Higgs-Audio V3 TTS + +> Multilingual text-to-speech with voice cloning on 1×H100 + +## Summary + +- Vendor: Boson AI +- Model: `bosonai/higgs-audio-v3-tts-4b` +- Task: Text-to-speech synthesis with optional voice cloning (100+ languages) +- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API; offline batch inference +- Maintainer: @yuekaiz + +## When to use this recipe + +Use this recipe to serve `bosonai/higgs-audio-v3-tts-4b` for high-quality +multilingual TTS. The model generates 24 kHz speech, supports zero-shot voice +cloning from a reference clip, and handles 100+ languages with inline control +tokens for emotion, style, and prosody. The architecture is a ~4B Qwen3 +backbone with fused multi-codebook embedding/head (8 codebooks × 1026 vocab, +MusicGen-style delay pattern). + +## References + +- Model card: [bosonai/higgs-audio-v3-tts-4b](https://huggingface.co/bosonai/higgs-audio-v3-tts-4b) +- Offline example: [`examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py`](../../examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py) +- Online example: [`examples/online_serving/text_to_speech/higgs_audio_v3/`](../../examples/online_serving/text_to_speech/higgs_audio_v3/) +- Benchmark results: see Performance section below + +## Hardware Support + +## GPU + +### 1×H100 80GB + +#### Environment + +- OS: Linux +- Python: 3.12+ +- CUDA: 12.x +- vLLM version: 0.22.0 +- vLLM-Omni version or commit: `36e048fd` (branch `higgs-v3`) + +#### Command + +**Online serving:** + +```bash +vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ + --host 0.0.0.0 --port 8095 \ + --trust-remote-code --omni +``` + +The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is +loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`). +It matches the high-throughput profile. + +**Offline batch inference:** + +```bash +python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \ + --texts "Hello world." "The quick brown fox jumps over the lazy dog." \ + --output-dir results/higgs_v3_wavs +``` + +**Offline voice clone:** + +```bash +python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \ + --texts "Text to synthesize in the cloned voice." \ + --ref-audio path/to/reference.wav \ + --ref-text "Transcript of the reference clip." \ + --output-dir results/higgs_v3_clone +``` + +#### Verification + +Basic TTS via curl: + +```bash +curl -X POST http://localhost:8095/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "bosonai/higgs-audio-v3-tts-4b", + "input": "Hello, how are you?" + }' --output hello.wav +``` + +Voice clone via Python client: + +```bash +python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \ + --base-url http://localhost:8095 \ + --model bosonai/higgs-audio-v3-tts-4b \ + --ref-audio path/to/reference.wav \ + --ref-text "Transcript of the reference." \ + --prompts "Text to clone." +``` + +#### Notes + +- Memory usage: Stage 0 (talker, ~4B) uses ~60% GPU memory; Stage 1 (codec decoder) uses ~25%. +- Key flags: `--trust-remote-code` and `--omni` are required. +- Output: 24 kHz mono WAV. +- Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity. +- Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`). + - `max_num_seqs=16` for both stages. + - Stage 0 and Stage 1 default to the same device (`0`) for single-GPU serving. + - Stage 0 intentionally keeps `enforce_eager=true`. This preserves the Higgs-specific local MLP CUDA graph path, which is the current high-throughput default. + - Stage 1 remains `enforce_eager=true` for the codec decoder. +- Deploy profiles: + - High throughput: `vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml`. + Use this for medium/high concurrency. The auto-discovered + `higgs_multimodal_qwen3.yaml` is kept as a compatibility/default alias for + this profile. + - Low latency: `vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml`. + Use this for low-concurrency serving (for example c1-c4). It sets Stage 0 + `enforce_eager=false` and enables vLLM `FULL_DECODE_ONLY` CUDA graph through + YAML `compilation_config`; no environment variable is required. + - Profile details: `vllm_omni/deploy/README_higgs_audio_v3.md`. +- Performance note: + - Do not switch the auto-discovered default Stage 0 profile to vLLM + `FULL_DECODE_ONLY` CUDA graph without an end-to-end throughput and + audio-quality revalidation. On the H20 SeedTTS c16/full-dataset benchmark, + the eager Stage 0 path with Higgs local MLP CUDA graph reproduced + ~35 audio_s/s (`1088/1088` OK per run, three runs around 134-136s). A + separate H20 low-concurrency smoke after the FULL_DECODE audio-feedback + capture fix showed lower c1/c4 latency, so the FULL_DECODE profile is kept + as an explicit low-latency option rather than the throughput default. +- Known limitations: + - Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture). + - Stage 0 full-decode CUDA graph is experimental; sampler, delay-state updates, staging, and request postprocess remain outside the graph. diff --git a/recipes/README.md b/recipes/README.md index 4648dbd6e..2bbc518df 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -33,11 +33,12 @@ recipes/ | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB | | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 | | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB | +| [`inclusionAI/Ming-omni-tts-0.5B.md`](./inclusionAI/Ming-omni-tts-0.5B.md) | Offline + online dense Ming TTS/audio generation | 1x AMD `gfx942` (ROCm 7.2) | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB | | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | -| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | -| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 | +| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | +| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 | | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB | | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB | | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 | @@ -50,6 +51,7 @@ recipes/ | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 | | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) | | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB | +| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB | | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB | | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB | | diff --git a/recipes/StabilityAI/Stable-Diffusion-3.5.md b/recipes/StabilityAI/Stable-Diffusion-3.5.md new file mode 100644 index 000000000..7c8f747aa --- /dev/null +++ b/recipes/StabilityAI/Stable-Diffusion-3.5.md @@ -0,0 +1,139 @@ +# Stable Diffusion 3.5 + +> Text-to-image serving with Stability AI's SD 3.5 family + +## Summary + +- Vendor: Stability AI +- Model: `stabilityai/stable-diffusion-3.5-medium`, `stabilityai/stable-diffusion-3.5-large` +- Task: Text-to-image generation +- Mode: Online serving with the OpenAI-compatible images API +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a lightweight, fast text-to-image model that +fits on a single consumer or workstation GPU. SD 3.5-medium requires only +~17 GiB VRAM and generates 1024×1024 images in ~4 seconds, making it one of +the most accessible diffusion models in vLLM-Omni. SD 3.5-large produces +noticeably sharper images at the cost of ~34 GiB VRAM and ~21 seconds per +image. + +## References + +- Model (medium): +- Model (large): +- Related example under `examples/`: + [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md) +- User guide: + [`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md) + +## Hardware Support + +This recipe documents two CUDA GPU configurations on a single RTX A6000 48 GB, +one for each model variant. Extend it with more hardware sections as community +validation lands. + +## GPU + +### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-medium + +#### Environment + +- OS: Ubuntu 22.04 +- Python: 3.12 +- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0 +- vLLM version: 0.22.0 +- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d + +#### Command + +```bash +vllm serve stabilityai/stable-diffusion-3.5-medium --omni --port 8091 +``` + +#### Verification + +```bash +curl -X POST http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "stabilityai/stable-diffusion-3.5-medium", + "prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus", + "negative_prompt": "blurry, out of focus", + "guidance_scale": 4.5, + "n": 1, + "size": "1024x1024" + }' | python3 -c " +import json, base64, sys +data = json.load(sys.stdin) +with open('output.png', 'wb') as f: + f.write(base64.b64decode(data['data'][0]['b64_json'])) +print('saved output.png') +" +``` + +Or use the example client: + +```bash +python examples/online_serving/text_to_image/openai_chat_client.py \ + --server http://localhost:8091 \ + --prompt "a cute cat sitting on a windowsill" \ + --output output.png +``` + +#### Notes + +- Memory usage: Model weights ~15.6 GiB, peak VRAM ~17.0 GiB (no CFG), ~20.9 GiB with `guidance_scale > 1` (CFG enabled). +- Generation time: ~4 seconds without CFG, ~8 seconds with `guidance_scale=4.5`. +- Model loading time: ~55 seconds (including weight download on first run). +- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG). + +### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-large + +#### Environment + +- OS: Ubuntu 22.04 +- Python: 3.12 +- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0 +- vLLM version: 0.22.0 +- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d + +#### Command + +```bash +vllm serve stabilityai/stable-diffusion-3.5-large --omni --port 8091 +``` + +> **Note**: `stable-diffusion-3.5-large` is a gated model. Accept the license +> on [Hugging Face](https://huggingface.co/stabilityai/stable-diffusion-3.5-large) +> and set `HF_TOKEN` before running. + +#### Verification + +```bash +curl -X POST http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "stabilityai/stable-diffusion-3.5-large", + "prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus", + "negative_prompt": "blurry, out of focus", + "guidance_scale": 4.5, + "n": 1, + "size": "1024x1024" + }' | python3 -c " +import json, base64, sys +data = json.load(sys.stdin) +with open('output.png', 'wb') as f: + f.write(base64.b64decode(data['data'][0]['b64_json'])) +print('saved output.png') +" +``` + +#### Notes + +- Memory usage: Peak VRAM ~31.6 GiB at 1024×1024, ~45.5 GiB at 2048×2048 with `guidance_scale=4.5`. +- Generation time: ~19 seconds at 1024×1024, ~102 seconds at 2048×2048. +- Image quality is noticeably better than medium, with sharper details. +- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG). +- 2048×2048 fits within 48 GB but leaves little headroom (~45.5 GiB used out of 49 GiB). diff --git a/recipes/Wan-AI/Wan2.2-S2V.md b/recipes/Wan-AI/Wan2.2-S2V.md index 9f34509d4..a91e768d0 100644 --- a/recipes/Wan-AI/Wan2.2-S2V.md +++ b/recipes/Wan-AI/Wan2.2-S2V.md @@ -171,6 +171,34 @@ python examples/offline_inference/speech_to_video/speech_to_video.py \ ``` +### Online Serving + +#### Server + +```bash +vllm serve Wan-AI/Wan2.2-S2V-14B --omni \ + --model-class-name WanS2VPipeline \ + --tensor-parallel-size 2 \ + --flow-shift 3.0 \ + --vae-use-slicing --vae-use-tiling \ + --port 8091 +``` + +#### Client + +```bash +no_proxy=127.0.0.1 \ +curl -X POST http://127.0.0.1:8091/v1/videos/sync \ + -F "prompt=A person singing" \ + -F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \ + -F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \ + -F "width=832" -F "height=480" \ + -F "num_inference_steps=40" \ + -F "guidance_scale=4.5" \ + -F "fps=16" \ + --output s2v_480p_serve.mp4 +``` + #### Verification The output video should show a talking/singing person matching the reference diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 57eb9048c..2ee9c2ed7 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -1,34 +1,43 @@ # Cosmos3-Nano -> Text-to-image, text-to-video, and image-to-video serving +> Text-to-image, text-to-video, image-to-video, and video-to-video serving ## Summary - Vendor: NVIDIA - Model: `nvidia/Cosmos3-Nano` -- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound), action policy +- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy - Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API - Maintainer: Community ## When to use this recipe Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation. -A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes; -the mode is selected per request: +A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves these modes; the +mode is selected per request: - **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`). -- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image. +- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image/video. - **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file upload, or `image_reference` JSON). +- **V2V** — `POST /v1/videos/sync` with a reference video (`input_reference` + video upload, or `video_reference` JSON). Cosmos3 conditions on selected + reference-video latent frames; use `extra_params.condition_frame_indexes_vision` + and `extra_params.condition_video_keep` to choose which prefix/tail frames guide + generation. - **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples. - **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks: - - `forward_dynamics` — given a first frame **and** an action trajectory, roll out - the resulting video. Synchronous: `POST /v1/videos/sync`. - - `policy` — given a first frame and a language instruction, **predict** the action - trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and - read the predicted action from the top-level `action` field + - `forward_dynamics` — given a first frame or video **and** an action trajectory, + roll out the resulting video. Synchronous: `POST /v1/videos/sync`. + - `policy` — given a first frame or video and a language instruction, + **predict** the action trajectory (and a rollout video). Use the async + `POST /v1/videos` endpoint and read the predicted action from the top-level + `action` field. + - `inverse_dynamics` — given a video, **recover** the action trajectory. Use + the async `POST /v1/videos` endpoint and read the recovered action from + the top-level `action` field (`{data, shape, dtype, raw_action_dim, domain_id}`). Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`, @@ -38,8 +47,9 @@ the mode is selected per request: **`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way (`domain_name=droid_lerobot`). - `inverse_dynamics` (recover the action from a given video) is supported by the - pipeline; **online inference of inverse dynamics will be added in a follow-up MR.** + Action requests can use `input_reference` or `video_reference` for video input. + `policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics` + requires a video reference. ## References @@ -142,6 +152,36 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \ -o cosmos3_i2v.mp4 +# Video-to-video -> /v1/videos/sync with an uploaded reference video. +# By default Cosmos3 conditions on latent indexes [0, 1]. For the default +# temporal VAE stride this decodes only the first 5 input frames. +# The model works best when the prompt describes the actual situation happening in the video. +# Generic prompts may create sub-standard generations. +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Nano" \ + -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \ + -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" \ + -F "num_inference_steps=35" -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true,"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \ + -F "seed=2222" \ + -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \ + -o cosmos3_v2v.mp4 + +# V2V can also use a JSON-safe URL/data-URL video reference. Do not combine +# video_reference with input_reference or image_reference. +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Nano" \ + -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" \ + -F "num_inference_steps=35" -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"last"}' \ + -F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \ + -o cosmos3_v2v_from_url.mp4 # Text-to-video-with-sound curl -sS -X POST http://localhost:8000/v1/videos/sync \ @@ -190,6 +230,21 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \ # poll until status == completed, then: curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}' curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4 + +# Action — inverse dynamics (video -> recovered action trajectory). +# Asynchronous: use the job metadata to read the recovered action. +VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \ + -H "Accept: application/json" \ + --form-string "model=nvidia/Cosmos3-Nano" \ + --form-string "prompt=Recover the robot action trajectory from this clip." \ + -F "input_reference=@motion_clip.mp4;type=video/mp4" \ + -F "size=640x480" -F "num_frames=17" -F "fps=5" \ + -F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \ + --form-string 'extra_params={"action_mode":"inverse_dynamics","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \ + -F "seed=0" | jq -r '.id') +# poll until status == completed, then: +curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}' +curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4 ``` #### Notes @@ -206,24 +261,27 @@ curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_polic - **Determinism:** identical seed reproduces identical output on the same hardware; outputs are not bit-identical across different GPU types. - **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1, - 3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720, - 189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`. + 3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V + 1280×720, 189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`. - **Key flags / params:** `--no-guardrails` (server) or `extra_params={"guardrails":false}` (per request) toggles safety. The per-request flag only takes effect when the server was launched **with** guardrails enabled (it cannot re-enable them on a `--no-guardrails` server). `use_resolution_template` / `use_duration_template` are off by default and only needed when not using upsampled prompts that already encode resolution/duration. + For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent + frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the + API decodes the first or last needed reference frames (`"first"` by default). - **Known limitations:** - Guardrails-on requires `cosmos-guardrail` **and** access to the gated `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise the server fails at pipeline build with a gated-repo / safety-checker error. - A guardrail-blocked prompt currently returns HTTP 500 (`"Guardrail blocked prompt"`). - - Action `forward_dynamics` (sync `/v1/videos/sync`) and `policy` (async - `/v1/videos`, returns the predicted action under the top-level `action` - field) are supported online. **Online inference of inverse dynamics will be - added in a follow-up MR.** + - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported + online. Use async `POST /v1/videos` when you need the predicted/recovered + action payload under the top-level `action` field; sync `/v1/videos/sync` + returns raw MP4 bytes and does not expose action metadata in the response body. ### 1x GPU (Offline generation) @@ -287,8 +345,9 @@ python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image - Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by `prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus - `num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video, - frames are returned in `outputs[0].request_output.images` as an + `num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and + `multi_modal_data={"video": []}` or a video tensor/array for V2V. + For video, frames are returned in `outputs[0].request_output.images` as an `(B, F, H, W, 3)` array. - The offline entry must be guarded by `if __name__ == "__main__":` — the engine spawns workers with the `spawn` start method. diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index 1dc88d0c3..7e63ba26d 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/recipes/cosmos3/Cosmos3-Super.md @@ -1,12 +1,12 @@ # Cosmos3-Super -> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio) +> Frontier 64B world model: text-to-image, text-to-video, image-to-video, video-to-video (+ optional audio) ## Summary - Vendor: NVIDIA - Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`) -- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound) +- Task: T2I, T2V, I2V, V2V generation, with optional synchronized audio (video + sound) - Mode: Online serving with the OpenAI-compatible image/video APIs - Maintainer: Community @@ -16,8 +16,8 @@ Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the recommended parallelism differ. Mode is selected per request (T2I → -`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true` -for audio). +`/v1/images/generations`; T2V/I2V/V2V → `/v1/videos/sync`; add +`generate_sound=true` for audio). ## References @@ -62,8 +62,9 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs. #### Verification Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full -T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24, -num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`. +T2I/T2V/I2V/V2V/T2VS curls); official params: `size=1280x720, num_frames=189, +fps=24, num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, +max_sequence_length=4096`. ```bash curl http://localhost:8000/v1/models @@ -84,6 +85,15 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ -F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \ -o cosmos3_super_i2v.mp4 +# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last". +curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \ + -F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \ + -o cosmos3_super_v2v.mp4 + # T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo) curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \ @@ -103,10 +113,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo** - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.) - **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk. -- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration` - semantics as Nano, including the **action** modality: `forward_dynamics` - (sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under - the top-level `action` field) — see the Cosmos3-Nano recipe for the request shape. - Online inference of inverse dynamics will be added in a follow-up MR. Verified on - the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted - action (`[16, 10]`) and the rollout video reliably. +- Same generation defaults, supported sizes, V2V reference-video controls + (`condition_frame_indexes_vision`, `condition_video_keep`), and + `generate_sound`/`sound_duration` + semantics as Nano, including the **action** modality: `forward_dynamics`, + `policy`, and `inverse_dynamics` — see the Cosmos3-Nano recipe for the request + shapes. Use async `/v1/videos` when you need predicted/recovered action metadata + under the top-level `action` field. Verified on the 64B Super under + `--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`) + and the rollout video reliably. diff --git a/recipes/inclusionAI/Ming-flash-omni-2.0.md b/recipes/inclusionAI/Ming-flash-omni-2.0.md index 3dcda2cbe..4a3a86669 100644 --- a/recipes/inclusionAI/Ming-flash-omni-2.0.md +++ b/recipes/inclusionAI/Ming-flash-omni-2.0.md @@ -57,12 +57,6 @@ Adjust `devices` in the YAML to match your hardware. #### Command -Thinker only (text output): - -```bash -vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 -``` - Thinker + talker (text and/or audio output): ```bash @@ -72,6 +66,15 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 \ --log-stats ``` +Thinker only (text-only output): + +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 \ + --omni \ + --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml \ + --port 8091 +``` + `--log-stats` is optional but recommended while validating the deployment. #### Verification diff --git a/recipes/inclusionAI/Ming-omni-tts-0.5B.md b/recipes/inclusionAI/Ming-omni-tts-0.5B.md new file mode 100644 index 000000000..cd1728034 --- /dev/null +++ b/recipes/inclusionAI/Ming-omni-tts-0.5B.md @@ -0,0 +1,82 @@ +# Ming-omni-tts 0.5B + +> Offline and online dense Ming text-to-speech on ROCm + +## Summary + +- Vendor: inclusionAI +- Model: `inclusionAI/Ming-omni-tts-0.5B` +- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls +- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference +- Maintainer: Community + +## References + +- [Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B) +- [Upstream repository](https://github.com/inclusionAI/Ming-omni-tts) +- [Offline example](../../examples/offline_inference/text_to_speech/ming_tts/) +- [Online example](../../examples/online_serving/text_to_speech/ming_tts/) + +## Hardware Support + +## ROCm + +### 1x AMD `gfx942` + +#### Environment + +- OS: Ubuntu 22.04.5 LTS, x86_64 +- Python: 3.12.13 +- ROCm / HIP: 7.2.53211 +- vLLM version: 0.22.0 +- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410` +- Docker image: `vllm/vllm-omni-rocm:v0.22.0` + +#### Command + +From the vLLM-Omni repository root: + +```bash +docker run --rm \ + --group-add=video \ + --ipc=host \ + --cap-add=SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --device /dev/kfd \ + --device /dev/dri \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + -v "$PWD":/app/vllm-omni \ + -w /app/vllm-omni \ + -e VLLM_ROCM_USE_AITER=0 \ + -p 8091:8091 \ + vllm/vllm-omni-rocm:v0.22.0 \ + --model inclusionAI/Ming-omni-tts-0.5B \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --omni \ + --port 8091 \ + --enforce-eager +``` + +#### Verification + +```bash +python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ + --text "我觉得社会企业同个人都有责任" \ + --instruction-json '{"方言":"广粤话"}' \ + --ref-audio /path/to/yue_prompt.wav \ + --max-new-tokens 200 \ + --output dialect.wav +``` + +`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text` +for the dialect case. + +## Notes + +- The official ROCm image includes the platform dependencies. +- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md) + for interactive and source-build workflows. +- The [offline example](../../examples/offline_inference/text_to_speech/ming_tts/) + can be run from an interactive container. +- The tested environment uses `--enforce-eager`. +- Output is mono 44.1 kHz audio. diff --git a/requirements/common.txt b/requirements/common.txt index e649450af..fb435218a 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -22,4 +22,3 @@ prettytable>=3.8.0 aenum==3.1.16 pyzmq>=25.0.0 janus>=1.0.0 -pydub diff --git a/requirements/xpu.txt b/requirements/xpu.txt index c7732b1e1..e59520dec 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -1,2 +1,3 @@ -r common.txt onnxruntime>=1.23.2 +auto-round-lib diff --git a/tests/benchmarks/test_accuracy_bench_utils.py b/tests/benchmarks/test_accuracy_bench_utils.py index 6ceebb11b..222b8c2c7 100644 --- a/tests/benchmarks/test_accuracy_bench_utils.py +++ b/tests/benchmarks/test_accuracy_bench_utils.py @@ -345,6 +345,14 @@ def test_parse_score_payload_handles_raw_json_and_delimited_json(): assert parse_score_payload(wrapped)["score"] == [6] +def test_parse_score_payload_handles_qwen_vl_nested_score_dicts(): + qwen_style = '{"score": [{"naturalness": 8, "artifact_free": 7}], "reasoning": "good quality"}' + assert parse_score_payload(qwen_style)["score"] == [8, 7] + + nested_values = '{"score": [{"score": 9}, {"score": 6}], "reasoning": "ok"}' + assert parse_score_payload(nested_values)["score"] == [9, 6] + + def test_summarize_gedit_generated_records_groups_by_task_and_language(): records = [] for group in GEDIT_GROUPS[:2]: diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index cf492e7cc..d94996014 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py @@ -20,6 +20,7 @@ from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreEventType from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin # isort: on @@ -142,3 +143,300 @@ class TestReplaceSessionWithStreamingUpdate: assert session.events assert session.events[-1].type == EngineCoreEventType.QUEUED + + +class _RealignSchedulerStub(OmniSchedulerMixin): + """Minimal scheduler surface for ``_realign_request_status_to_queues``. + + Pins ``self.requests`` / ``self.running`` / ``self.waiting`` to the + same shapes the upstream ``Scheduler.finish_requests`` reads, so the + helper sees realistic state without spinning a real scheduler. + """ + + def __init__( + self, + *, + requests: dict, + running: list, + waiting: list, + ) -> None: + self.requests = requests + self.running = running + self.waiting = waiting + + +class TestRealignRequestStatusToQueues: + """Regression for the residual hang described in + https://github.com/vllm-project/vllm-omni/pull/3774 -- chunk-transfer + adapter's ``requests_origin_status`` table goes stale on the + ``waiting → running`` admit transition, and an abort that lands + before the next deque round-trip stomps stale ``WAITING`` onto a + request that actually lives in ``self.running``. After + ``max_num_seqs`` such aborts every ``input_batch`` slot is leaked and + new requests hang at ``chunks=0``. + """ + + def test_running_with_stale_waiting_status_is_realigned_to_running( + self, + ) -> None: + """The exact race the hang reproduces: a request lives in + ``self.running`` but its ``status`` is ``WAITING``. After + realign, status must be ``RUNNING`` so upstream + ``Scheduler.finish_requests`` removes it from ``self.running``. + """ + req = _make_request(request_id="req-stale") + req.status = RequestStatus.WAITING # stale -- actually in running + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[req], + waiting=[], + ) + sched._realign_request_status_to_queues([req.request_id]) + + assert req.status == RequestStatus.RUNNING + + def test_waiting_with_stale_running_status_is_realigned_to_waiting( + self, + ) -> None: + """Symmetric case: request lives in ``self.waiting`` but status + is ``RUNNING``. Upstream's else branch should run, so realign to + ``WAITING``. + """ + req = _make_request(request_id="req-stale-2") + req.status = RequestStatus.RUNNING # stale -- actually in waiting + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[], + waiting=[req], + ) + sched._realign_request_status_to_queues([req.request_id]) + + assert req.status == RequestStatus.WAITING + + def test_already_aligned_status_is_left_unchanged(self) -> None: + """Healthy case: status matches actual queue. No status mutation + means no spurious side effects. + """ + req_running = _make_request(request_id="req-r") + req_running.status = RequestStatus.RUNNING + req_waiting = _make_request(request_id="req-w") + req_waiting.status = RequestStatus.WAITING + + sched = _RealignSchedulerStub( + requests={ + req_running.request_id: req_running, + req_waiting.request_id: req_waiting, + }, + running=[req_running], + waiting=[req_waiting], + ) + sched._realign_request_status_to_queues([req_running.request_id, req_waiting.request_id]) + + assert req_running.status == RequestStatus.RUNNING + assert req_waiting.status == RequestStatus.WAITING + + def test_unknown_request_id_is_skipped_silently(self) -> None: + """Aligning ids that aren't in ``self.requests`` (already freed + upstream, or never existed) must be a no-op. + """ + sched = _RealignSchedulerStub(requests={}, running=[], waiting=[]) + sched._realign_request_status_to_queues(["never-existed"]) # no raise + + def test_request_in_neither_queue_is_left_unchanged(self) -> None: + """If a tracked request is in neither queue (e.g. parked in a + chunk-transfer deque), realign must not invent a status. The + adapter / deque purge owns that surface; the realign helper is + only for the admit-transition staleness. + """ + req = _make_request(request_id="req-parked") + req.status = RequestStatus.WAITING_FOR_CHUNK + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[], + waiting=[], + ) + sched._realign_request_status_to_queues([req.request_id]) + + assert req.status == RequestStatus.WAITING_FOR_CHUNK + + def test_request_ids_str_is_treated_as_single_id(self) -> None: + """Match upstream ``Scheduler.finish_requests`` resolution: a + bare string is treated as one id, not iterated as characters. + """ + req = _make_request(request_id="req-s") + req.status = RequestStatus.WAITING + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[req], + waiting=[], + ) + sched._realign_request_status_to_queues(req.request_id) + + assert req.status == RequestStatus.RUNNING + + def test_request_ids_none_aligns_every_known_request(self) -> None: + """``request_ids=None`` matches upstream's "all requests" path. + Realign must walk every entry in ``self.requests`` and fix any + stale status it finds. + """ + stale = _make_request(request_id="req-stale-none") + stale.status = RequestStatus.WAITING # but actually in running + + clean = _make_request(request_id="req-clean-none") + clean.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={ + stale.request_id: stale, + clean.request_id: clean, + }, + running=[stale, clean], + waiting=[], + ) + sched._realign_request_status_to_queues(None) + + assert stale.status == RequestStatus.RUNNING + assert clean.status == RequestStatus.RUNNING + + def test_finished_request_is_skipped(self) -> None: + """Already-finished requests must not be touched -- they may + have legitimate finished statuses (FINISHED_STOPPED etc.) that + a status flip would corrupt. + """ + req = _make_request(request_id="req-finished") + req.status = RequestStatus.FINISHED_STOPPED + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[req], # not realistic, but we want to prove the guard fires + waiting=[], + ) + sched._realign_request_status_to_queues([req.request_id]) + + assert req.status == RequestStatus.FINISHED_STOPPED + + +class TestPurgeFinishedFromRunning: + """Regression for the residual ``self.running`` slot leak surface + paired with ``_realign_request_status_to_queues``: even after + realign + ``super().finish_requests`` runs, corner cases can leave + already-finished or untracked entries in ``self.running`` -- e.g. + a connector cleanup that pops from ``self.requests`` without + unwinding ``self.running``, or a ``status`` mid-transition when + finish ran. The defensive post-finish purge sweeps those residues + so the worker's ``input_batch`` slot never pins a freed request. + + See https://github.com/vllm-project/vllm-omni/pull/3774 discussion + and the residual-hang reproduction in that PR. + """ + + def test_finished_request_is_purged_from_running(self) -> None: + """``is_finished()`` request lingering in ``self.running`` must + be swept so its ``input_batch`` slot is freed.""" + finished = _make_request(request_id="req-finished") + finished.status = RequestStatus.FINISHED_STOPPED # is_finished() True + + sched = _RealignSchedulerStub( + requests={finished.request_id: finished}, + running=[finished], + waiting=[], + ) + sched._purge_finished_from_running() + + assert sched.running == [] + + def test_untracked_request_is_purged_from_running(self) -> None: + """Request lingering in ``self.running`` but no longer present + in ``self.requests`` (already freed by upstream / connector) + must be swept.""" + untracked = _make_request(request_id="req-untracked") + untracked.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={}, # already deleted from self.requests + running=[untracked], + waiting=[], + ) + sched._purge_finished_from_running() + + assert sched.running == [] + + def test_healthy_running_is_left_unchanged(self) -> None: + """Live tracked running requests must not be touched -- the + purge is defensive, not aggressive.""" + alive = _make_request(request_id="req-alive") + alive.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={alive.request_id: alive}, + running=[alive], + waiting=[], + ) + sched._purge_finished_from_running() + + assert sched.running == [alive] + + def test_empty_running_is_noop(self) -> None: + """Empty ``self.running`` must short-circuit cleanly.""" + sched = _RealignSchedulerStub(requests={}, running=[], waiting=[]) + sched._purge_finished_from_running() + + assert sched.running == [] + + def test_mixed_alive_and_dead_keeps_only_alive_in_order(self) -> None: + """Mixed ``self.running`` -- some alive, some finished, some + untracked. Sweep keeps only the alive ones and preserves their + relative order.""" + alive_a = _make_request(request_id="req-alive-a") + alive_a.status = RequestStatus.RUNNING + finished_b = _make_request(request_id="req-finished-b") + finished_b.status = RequestStatus.FINISHED_STOPPED + untracked_c = _make_request(request_id="req-untracked-c") + untracked_c.status = RequestStatus.RUNNING + alive_d = _make_request(request_id="req-alive-d") + alive_d.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={ + alive_a.request_id: alive_a, + finished_b.request_id: finished_b, # tracked but is_finished + # untracked_c is NOT in self.requests + alive_d.request_id: alive_d, + }, + running=[alive_a, finished_b, untracked_c, alive_d], + waiting=[], + ) + sched._purge_finished_from_running() + + assert sched.running == [alive_a, alive_d] + + def test_in_place_mutation_keeps_helper_local_list_identity(self) -> None: + """The slice-assign form (``self.running[:] = ...``) avoids an + extra rebind inside this helper. Note that across a full + ``finish_requests`` call upstream ``Scheduler.finish_requests`` + already rebinds ``self.running``, so global identity is not a + contract -- this test only pins the local helper behavior so a + future refactor away from in-place mutation is intentional. + """ + finished = _make_request(request_id="req-finished") + finished.status = RequestStatus.FINISHED_STOPPED + alive = _make_request(request_id="req-alive") + alive.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={ + finished.request_id: finished, + alive.request_id: alive, + }, + running=[finished, alive], + waiting=[], + ) + held_ref = sched.running + sched._purge_finished_from_running() + + assert held_ref is sched.running + assert held_ref == [alive] diff --git a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py new file mode 100644 index 000000000..9b564d7d5 --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py @@ -0,0 +1,158 @@ +"""Integration tests for the post-finish defensive purge of +``self.running`` placed in +``OmniARScheduler.finish_requests`` / +``OmniGenerationScheduler.finish_requests``. + +Pins *placement* — that ``_purge_finished_from_running`` runs AFTER +``super().finish_requests`` so it catches entries upstream's removal +missed (e.g. status mid-transition, connector cleanup popping +``self.requests`` without unwinding ``self.running``). Direct helper +unit tests live in ``test_omni_scheduler_mixin.py``; these tests +exercise the full ``finish_requests`` path end-to-end. + +Regression for the residual self.running slot leak surface paired +with ``_realign_request_status_to_queues`` (vllm-omni#3774). +""" + +from __future__ import annotations + +import pytest +from vllm.v1.request import RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod +import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _StubRequest: + """Minimal Request stub with the surface the helpers exercise.""" + + def __init__(self, request_id: str, status: RequestStatus) -> None: + self.request_id = request_id + self.status = status + + def is_finished(self) -> bool: + return RequestStatus.is_finished(self.status) + + +_SCHEDULER_PARAMS = [ + pytest.param(OmniARScheduler, ar_sched_mod, id="ar"), + pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"), +] + + +def _make_scheduler(scheduler_cls, *, requests, running, waiting): + """Bare scheduler with just the surface ``finish_requests`` reads.""" + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler.requests = requests + scheduler.running = running + scheduler.waiting = waiting + return scheduler + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_purges_finished_request_from_running( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + """A request whose ``status`` is a finished value but still lives + in ``self.running`` after ``super().finish_requests`` returns must + be swept by the post-finish purge. Mocks ``super()`` to simulate + the corner case where upstream's status-driven removal missed the + entry. + """ + finished = _StubRequest("req-finished", RequestStatus.FINISHED_STOPPED) + scheduler = _make_scheduler( + scheduler_cls, + requests={finished.request_id: finished}, + running=[finished], + waiting=[], + ) + + def fake_finish_requests(self, request_ids, finished_status): + # super() returning the abort tuple but leaving self.running untouched + return [(finished.request_id, 0)] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + result = scheduler_cls.finish_requests( + scheduler, + [finished.request_id], + RequestStatus.FINISHED_STOPPED, + ) + + assert scheduler.running == [] + assert result == [(finished.request_id, 0)] + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_purges_untracked_request_from_running( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + """A request that has been deleted from ``self.requests`` (e.g. + after upstream's ``_free_request`` ran but a status mismatch + caused the ``self.running`` removal branch to no-op) must be + swept by the post-finish purge. + """ + untracked = _StubRequest("req-untracked", RequestStatus.RUNNING) + scheduler = _make_scheduler( + scheduler_cls, + requests={}, # already deleted from self.requests + running=[untracked], + waiting=[], + ) + + def fake_finish_requests(self, request_ids, finished_status): + return [(untracked.request_id, 0)] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + result = scheduler_cls.finish_requests( + scheduler, + [untracked.request_id], + RequestStatus.FINISHED_STOPPED, + ) + + assert scheduler.running == [] + assert result == [(untracked.request_id, 0)] + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_leaves_healthy_running_intact( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + """A live ``RUNNING`` request that is not in the finish set must + not be swept by the defensive purge -- this pins that the predicate + is conservative and won't aggressively drop healthy entries. + """ + alive = _StubRequest("req-alive", RequestStatus.RUNNING) + leaving_id = "req-leaving" + scheduler = _make_scheduler( + scheduler_cls, + requests={alive.request_id: alive}, + running=[alive], # leaving already removed by super() (typical happy path) + waiting=[], + ) + + def fake_finish_requests(self, request_ids, finished_status): + return [(leaving_id, 0)] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + scheduler_cls.finish_requests( + scheduler, + [leaving_id], + RequestStatus.FINISHED_STOPPED, + ) + + assert scheduler.running == [alive] diff --git a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py index 4088cbf78..fcac8172b 100644 --- a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py +++ b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py @@ -46,6 +46,9 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids( scheduler = scheduler_cls.__new__(scheduler_cls) scheduler.chunk_transfer_adapter = None scheduler.input_coordinator = coordinator + scheduler.requests = {} + scheduler.running = [] + scheduler.waiting = [] def fake_finish_requests(self, request_ids, finished_status): assert request_ids == ["req-a", "req-b"] diff --git a/tests/dfx/perf/tests/test_higgs_audio_v3.json b/tests/dfx/perf/tests/test_higgs_audio_v3.json new file mode 100644 index 000000000..6fdaec580 --- /dev/null +++ b/tests/dfx/perf/tests/test_higgs_audio_v3.json @@ -0,0 +1,74 @@ +[ + { + "test_name": "test_higgs_audio_v3", + "server_params": { + "model": "bosonai/higgs-audio-v3-tts-4b", + "stage_overrides": { + "0": { + "gpu_memory_utilization": 0.3 + }, + "1": { + "gpu_memory_utilization": 0.2 + } + }, + "extra_cli_args": [ + "--deploy-config", + "vllm_omni/deploy/higgs_multimodal_qwen3.yaml", + "--trust-remote-code" + ] + }, + "benchmark_params": [ + { + "task": "default_voice", + "eval_phase": "latency", + "dataset_name": "seed-tts-text", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", + "num_prompts": [ + 20 + ], + "max_concurrency": [ + 1 + ], + "seed_tts_locale": "en", + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "median_audio_ttfp_ms": [ + 3000 + ], + "median_audio_rtf": [ + 0.5 + ] + } + }, + { + "task": "default_voice", + "eval_phase": "throughput", + "dataset_name": "seed-tts-text", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", + "num_prompts": [ + 500 + ], + "max_concurrency": [ + 16 + ], + "seed_tts_locale": "en", + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "median_audio_ttfp_ms": [ + 5000 + ], + "median_audio_rtf": [ + 1.2 + ], + "audio_throughput": [ + 30 + ] + } + } + ] + } +] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json index 0e7cb63f7..bee799a77 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json @@ -23,10 +23,10 @@ "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [1000, 2000, 3000, 4000], - "mean_tpot_ms": [200, 200, 200, 200], - "mean_audio_ttfp_ms": [10000, 15000, 20000, 25000], - "mean_audio_rtf": [1.5, 1.5, 1.5, 1.5] + "mean_ttft_ms": [203, 283, 470, 869], + "mean_tpot_ms": [28, 36, 41, 45], + "mean_audio_ttfp_ms": [861, 1200, 1528, 2100], + "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39] } }, { @@ -49,9 +49,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [2000], - "mean_audio_ttfp_ms": [2000], - "mean_audio_rtf": [0.25] + "mean_ttft_ms": [377], + "mean_audio_ttfp_ms": [684], + "mean_audio_rtf": [0.14] } }, { @@ -76,9 +76,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [6000], - "mean_audio_ttfp_ms": [6000], - "mean_audio_rtf": [0.7] + "mean_ttft_ms": [228], + "mean_audio_ttfp_ms": [686], + "mean_audio_rtf": [0.16] } }, { @@ -105,9 +105,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [12000], - "mean_audio_ttfp_ms": [12000], - "mean_audio_rtf": [1.0] + "mean_ttft_ms": [522], + "mean_audio_ttfp_ms": [1119], + "mean_audio_rtf": [0.21] } } ] diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json index 733c40963..9246b502f 100644 --- a/tests/dfx/perf/tests/test_tts.json +++ b/tests/dfx/perf/tests/test_tts.json @@ -339,5 +339,66 @@ } } ] + }, + { + "test_name": "test_voxtral_tts", + "server_params": { + "model": "mistralai/Voxtral-4B-TTS-2603" + }, + "benchmark_params": [ + { + "task": "default_voice", + "eval_phase": "latency", + "dataset_name": "seed-tts-text", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", + "num_prompts": [20], + "max_concurrency": [1], + "seed_tts_locale": "en", + "extra_body": {"voice": "casual_male"}, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "median_audio_ttfp_ms": [150], + "median_audio_rtf": [0.15] + } + }, + { + "task": "default_voice", + "eval_phase": "throughput", + "dataset_name": "seed-tts-text", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", + "num_prompts": [80], + "max_concurrency": [8], + "seed_tts_locale": "en", + "extra_body": {"voice": "casual_male"}, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "median_audio_ttfp_ms": [1500], + "median_audio_rtf": [0.30], + "audio_throughput": [30.0] + } + }, + { + "task": "default_voice", + "eval_phase": "stress", + "dataset_name": "seed-tts-text", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", + "num_prompts": [100], + "request_rate": [2.0], + "seed_tts_locale": "en", + "extra_body": {"voice": "casual_male"}, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "median_audio_ttfp_ms": [3000], + "median_audio_rtf": [0.25], + "audio_throughput": [4.0] + } + } + ] } ] diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py index 9d1782e76..9e39b059b 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py @@ -198,19 +198,78 @@ def test_video_chat_stream_invalid_requests( # WS /v1/realtime # ───────────────────────────────────────────────────────────────────────────── +# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake). +_REALTIME_WS_MODEL_MISMATCH = object() +_REALTIME_WS_INVALID_AUDIO_APPEND = object() + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize( + "send_frames_spec, ws_error_code, err_message", + [ + pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"), + pytest.param( + json.dumps({"type": "session.update"}), + "invalid_event", + ("model", "Missing required field"), + id="session_update_missing_model", + ), + pytest.param( + _REALTIME_WS_MODEL_MISMATCH, + "model_not_found", + ("model", "does not exist"), + id="session_update_model_mismatch", + ), + pytest.param( + json.dumps({"type": "input_audio_buffer.commit", "final": False}), + "model_not_validated", + ("session.update", "validate the model"), + id="commit_without_model_validation", + ), + pytest.param( + json.dumps({"type": "unknown.event"}), + "unknown_event", + "Unknown event type", + id="unknown_event_type", + ), + pytest.param( + _REALTIME_WS_INVALID_AUDIO_APPEND, + "invalid_audio", + "Invalid audio data", + id="invalid_audio_append", + ), + ], +) @pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True) -def test_realtime_rejected_when_async_chunk_enabled( +def test_realtime_invalid_requests( omni_server: OmniServer, openai_client: OpenAIClientHandler, + send_frames_spec: Any, + ws_error_code: str, + err_message: str | tuple[str, ...], ) -> None: + if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH: + send_frames = json.dumps( + { + "type": "session.update", + "model": "this-model-is-not-served-on-this-instance", + } + ) + elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND: + send_frames = [ + json.dumps({"type": "session.update", "model": omni_server.model}), + json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}), + ] + else: + send_frames = send_frames_spec openai_client.send_realtime_ws_request( { + "send_frames": send_frames, "timeout": 120, "ws_max_size": None, + "ws_skip_types": ["session.created"], "ws_json_type": "error", - "ws_error_code": "unsupported", - "err_message": "async_chunk", + "ws_error_code": ws_error_code, + "err_message": err_message, } ) diff --git a/tests/dfx/reliability/test_reliability_qwen3_omni.py b/tests/dfx/reliability/test_reliability_qwen3_omni.py index 49028a03f..28e8dedf7 100644 --- a/tests/dfx/reliability/test_reliability_qwen3_omni.py +++ b/tests/dfx/reliability/test_reliability_qwen3_omni.py @@ -319,6 +319,7 @@ QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLO @pytest.mark.slow +@pytest.mark.skip(reason="issue#4285") @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech( @@ -387,6 +388,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech( @pytest.mark.slow +@pytest.mark.skip(reason="issue#4285") @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None: @@ -429,6 +431,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi @pytest.mark.slow +@pytest.mark.skip(reason="issue#4285") @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None: @@ -771,7 +774,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( @pytest.mark.slow -@pytest.mark.skip(reason="OOM injection causes the serving instance to exit") +@pytest.mark.skip(reason="issue#4285") @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_state_converges_after_fault_removed( diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 9434702f9..376dbe685 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -1,19 +1,46 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for the DiffusersPipelineLoader. +""" + from types import SimpleNamespace import pytest import torch import torch.nn as nn +from huggingface_hub import snapshot_download +from vllm.config.load import LoadConfig from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter +from vllm_omni.diffusion.models.helios import HeliosPipeline from vllm_omni.diffusion.registry import initialize_model pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +model_path = "hf-internal-testing/tiny-helios-modular-pipe" + + +@pytest.fixture(scope="module") +def prefetch_helios_model(): + """Downloads the tiny helios model prior to running a test.""" + snapshot_download(model_path) + + +@pytest.fixture(scope="function") +def mock_tp_group(mocker): + """Mocks the tensor parallel group; this is needed to initialize the Helios model.""" + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1) + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0) + mock_group = mocker.MagicMock() + mock_group.world_size = 1 + mock_group.rank_in_group = 0 + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group) + class _DummyPipelineModel(nn.Module): def __init__(self, *, source_prefix: str): @@ -42,7 +69,13 @@ class _DummyPipelineModel(nn.Module): def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader: - loader = object.__new__(DiffusersPipelineLoader) + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False), + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader.counter_before_loading_weights = 0.0 loader.counter_after_loading_weights = 0.0 @@ -150,16 +183,15 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch): quantization_config=None, ) - loader = object.__new__(DiffusersPipelineLoader) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) loader.load_weights = lambda model: None # type: ignore[assignment] loader._process_weights_after_loading = lambda model, target_device: None # type: ignore[assignment] - loader._is_gguf_quantization = lambda _od_config: False # type: ignore[assignment] + loader._is_gguf_quantization = lambda: False # type: ignore[assignment] monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel) monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu")) model = loader.load_model( - od_config, load_device="cpu", load_format="custom_pipeline", custom_pipeline_name="tests.dummy.ConfigAwarePipeline", @@ -168,3 +200,33 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch): assert model.captured_config is od_config assert model.seen_config_during_init is od_config assert get_current_diffusion_config_or_none() is None + + +def test_get_all_weights(prefetch_helios_model, mock_tp_group): + """Ensure that get all weights on a tiny model resolves to nonempty weights.""" + od_config = OmniDiffusionConfig( + model_class_name="HeliosPipeline", + model=model_path, + ) + loader = DiffusersPipelineLoader( + load_config=LoadConfig(), + od_config=od_config, + ) + pipeline = HeliosPipeline(od_config=od_config) + + weights = list(loader.get_all_weights(pipeline)) + assert len(weights) > 0 + + +def test_load_model(prefetch_helios_model, mock_tp_group): + """Ensure that load model creates an instance of the expected pipeline class.""" + od_config = OmniDiffusionConfig( + model_class_name="HeliosPipeline", + model=model_path, + ) + loader = DiffusersPipelineLoader( + load_config=LoadConfig(), + od_config=od_config, + ) + model = loader.load_model(load_device="cpu") + assert isinstance(model, HeliosPipeline) diff --git a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py b/tests/diffusion/model_loader/test_diffusers_loader_gguf.py index 9aae9135e..f747dca81 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py +++ b/tests/diffusion/model_loader/test_diffusers_loader_gguf.py @@ -62,6 +62,7 @@ def _make_loader() -> DiffusersPipelineLoader: ignore_patterns=["*.tmp"], ) loader.od_config = None + loader.quant_config = None loader.counter_before_loading_weights = 0.0 loader.counter_after_loading_weights = 0.0 return loader @@ -69,53 +70,47 @@ def _make_loader() -> DiffusersPipelineLoader: def test_is_gguf_quantization_accepts_dict_config(): loader = _make_loader() - od_config = SimpleNamespace( - quantization_config={ - "method": "gguf", - "gguf_model": "weights.gguf", - } - ) + loader.quant_config = { + "method": "gguf", + "gguf_model": "weights.gguf", + } - assert loader._is_gguf_quantization(od_config) is True + assert loader._is_gguf_quantization() is True def test_is_gguf_quantization_rejects_non_gguf_dict_config(): loader = _make_loader() - od_config = SimpleNamespace( - quantization_config={ - "method": "fp8", - "gguf_model": "weights.gguf", - } - ) + loader.quant_config = { + "method": "fp8", + "gguf_model": "weights.gguf", + } - assert loader._is_gguf_quantization(od_config) is False + assert loader._is_gguf_quantization() is False def test_is_gguf_quantization_requires_gguf_model_for_dict_config(): loader = _make_loader() - od_config = SimpleNamespace(quantization_config={"method": "gguf"}) + loader.quant_config = {"method": "gguf"} with pytest.raises(ValueError, match="gguf_model"): - loader._is_gguf_quantization(od_config) + loader._is_gguf_quantization() def test_is_gguf_quantization_accepts_object_config(): loader = _make_loader() - od_config = SimpleNamespace( - quantization_config=SimpleNamespace( - get_name=lambda: "gguf", - gguf_model="weights.gguf", - ) + loader.quant_config = SimpleNamespace( + get_name=lambda: "gguf", + gguf_model="weights.gguf", ) - assert loader._is_gguf_quantization(od_config) is True + assert loader._is_gguf_quantization() is True def test_is_gguf_quantization_uses_fallback_object_without_get_name(): loader = _make_loader() - od_config = SimpleNamespace(quantization_config=SimpleNamespace(gguf_model="weights.gguf")) + loader.quant_config = SimpleNamespace(gguf_model="weights.gguf") - assert loader._is_gguf_quantization(od_config) is True + assert loader._is_gguf_quantization() is True def test_get_model_loadable_names_collects_parameters_and_buffers(): @@ -171,16 +166,14 @@ def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.Mon def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch): loader = _make_loader() + loader.quant_config = SimpleNamespace(gguf_model="weights.gguf") + loader.od_config = SimpleNamespace(revision="main") source = DiffusersPipelineLoader.ComponentSource( model_or_path="dummy", subfolder="transformer", revision=None, prefix="transformer.", ) - od_config = SimpleNamespace( - revision="main", - quantization_config=SimpleNamespace(gguf_model="weights.gguf"), - ) class _Adapter: def weights_iterator(self): @@ -190,7 +183,7 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf") monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter()) - weights = list(loader._get_gguf_weights_iterator(source, object(), od_config)) + weights = list(loader._get_gguf_weights_iterator(source, object())) assert weights[0][0] == "transformer.weight" assert torch.equal(weights[0][1], torch.ones((2, 2))) @@ -201,13 +194,12 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch): loader = _make_loader() model = _DummyModel() - od_config = SimpleNamespace() monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources)) monkeypatch.setattr( loader, "_get_gguf_weights_iterator", - lambda source, model, od_config: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]), + lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]), ) def _hf_iter(source): @@ -222,7 +214,7 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights( lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"}, ) - loaded = loader._load_weights_with_gguf(model, od_config) + loaded = loader._load_weights_with_gguf(model) assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"} assert model.calls[0] == ["transformer.weight"] @@ -233,13 +225,12 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights( def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch): loader = _make_loader() model = _DummyModel() - od_config = SimpleNamespace() monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources)) monkeypatch.setattr( loader, "_get_gguf_weights_iterator", - lambda source, model, od_config: iter( + lambda source, model: iter( [ ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)), ("transformer.bias", torch.zeros(2, dtype=torch.float32)), @@ -262,7 +253,7 @@ def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"}, ) - loaded = loader._load_weights_with_gguf(model, od_config) + loaded = loader._load_weights_with_gguf(model) assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"} assert hf_calls == ["vae"] diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 087d1a673..c51f3ff08 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -374,6 +374,19 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None: assert tuple(additional["preprocessed_image"].shape) == (1, 3, 16, 32) assert tuple(additional["preprocessed_video"].shape) == (1, 3, 3, 16, 32) + frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "purple", "black")] + v2v = SimpleNamespace( + prompts=[{"prompt": "Continue.", "multi_modal_data": {"video": frames}}], + sampling_params=SimpleNamespace( + height=16, + width=32, + extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"}, + ), + ) + additional = preprocess(v2v).prompts[0]["additional_information"] + assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32) + assert additional["condition_frame_indexes_vision"] == [0, 1] + def test_postprocess_handles_image_video_audio_and_validation() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func @@ -464,6 +477,19 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline) assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]] assert image_latent.shape == (1, 2, 1, 2, 3) + pipeline._encode_video_tensor = lambda *args, **kwargs: torch.full((1, 2, 3, 2, 3), 6.0) + v2v_latents, v2v_velocity_mask, v2v_condition = pipeline._prepare_latents_v2v( + torch.zeros(1, 3, 5, 16, 24), + 16, + 24, + 9, + torch.Generator(device="cpu").manual_seed(0), + [0, 1], + ) + torch.testing.assert_close(v2v_latents[:, :, 0:2], torch.full((1, 2, 2, 2, 3), 6.0)) + assert v2v_velocity_mask.tolist() == [[[[[0.0]], [[0.0]], [[1.0]]]]] + assert v2v_condition.shape == (1, 2, 3, 2, 3) + pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3) pipeline._sound_tokenizer = SimpleNamespace( sample_rate=10, @@ -670,6 +696,34 @@ class TestForwardRouting: ) assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask + video_tensor = torch.zeros(1, 3, 5, 16, 16) + v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0) + v2v_mask = torch.tensor([[[[[0.0]], [[1.0]]]]]) + pipeline._prepare_latents_v2v = lambda *args, **kwargs: ( + torch.zeros(1, 2, 2, 1, 1), + v2v_mask, + v2v_condition, + ) + pipeline.forward( + SimpleNamespace( + prompts=[ + { + "prompt": "continue", + "modalities": ["video"], + "additional_information": { + "preprocessed_video": video_tensor, + "condition_frame_indexes_vision": [0], + }, + } + ], + sampling_params=make_sampling_params(height=16, width=16, num_frames=5), + ) + ) + assert captured["flow_shifts"][-1] == 10.0 + assert captured["format"]["negative_prompt"] == "" + assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask + assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition + pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3) sound_latents = torch.zeros(1, 3, 4) pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 50a8d22ad..62411fc67 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -131,6 +131,38 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No assert tuple(output.shape) == (1, 2, 1, 2, 2) +def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + class RecordingGather(nn.Module): + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + def forward(self, hidden_gen: torch.Tensor) -> torch.Tensor: + self.calls += 1 + return hidden_gen + + monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) + model = transformer_cosmos3.Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) + ) + gather = RecordingGather() + model.gen_sp_gather = gather + + output = model( + hidden_states=torch.zeros(1, 2, 1, 2, 2), + timestep=torch.tensor([1.0]), + text_ids=torch.tensor([[1, 2]], dtype=torch.long), + text_mask=torch.ones(1, 2, dtype=torch.long), + video_shape=(1, 2, 2), + fps=24.0, + ) + + assert gather.calls == 1 + assert tuple(output.shape) == (1, 2, 1, 2, 2) + + def test_sound_and_action_modules_follow_config() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer diff --git a/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py b/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py new file mode 100644 index 000000000..dc4d49b11 --- /dev/null +++ b/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py @@ -0,0 +1,728 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for multi-resolution support in HunyuanImage3. + +Covers three areas: + +1. ``ResolutionGroup._calc_by_step`` — training-bucket generation, + verifying that each resolution has correct width×height and that + indices match expected ratio ordering. + +2. ``HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS`` — extra resolution matching + via ``Resolution.append`` and new-entry insertion in + ``ResolutionGroup.__init__``. + +3. ``get_cached_resolution_group`` and ``ar2diffusion`` — caching + behaviour and correct ``ResolutionGroup`` construction for the + AR→DiT stage transition. +""" + +from __future__ import annotations + +import math +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ( + HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS, + Resolution, + ResolutionGroup, + get_cached_resolution_group, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _make_reso(desc: str | tuple[int, int]) -> Resolution: + """Build a ``Resolution`` from a ``"HxW"`` string or ``(h, w)`` tuple.""" + if isinstance(desc, str): + h, w = desc.split("x") + return Resolution(int(h), int(w)) + return Resolution(desc[0], desc[1]) + + +def _reso_set(rg: ResolutionGroup) -> set[tuple[int, int]]: + """Return all ``(h, w)`` tuples reachable through a ``ResolutionGroup``. + + For entries that carry ``extra_res`` this includes every alternative + resolution in that set. + """ + out: set[tuple[int, int]] = set() + for r in rg.data: + out.add((r.h, r.w)) + out |= r.extra_res + return out + + +# --------------------------------------------------------------------------- +# 1. _calc_by_step — training-bucket resolution parsing +# --------------------------------------------------------------------------- + +_BASE = 1024 +_STEP = _BASE // 16 # 64 +_ALIGN = 16 +_MIN_SIDE = _BASE // 2 # 512 +_MAX_SIDE = _BASE * 2 # 2048 + + +def _build_plain_group( + base_size: int = _BASE, + step: int | None = None, + align: int = _ALIGN, +) -> ResolutionGroup: + """Return a ``ResolutionGroup`` with *no* extra resolutions.""" + return ResolutionGroup(base_size=base_size, step=step, align=align, extra_resolutions=None) + + +class TestCalcByStep: + """Verify ``_calc_by_step`` generates the expected training buckets.""" + + # -- count & range ------------------------------------------------------- + + def test_count_for_base_1024(self): + """33 buckets: 1 centre + 2×16 from the stepping loops.""" + rg = _build_plain_group() + # 1 centre + 16 portrait + 16 landscape + assert len(rg.data) == 33, f"expected 33 training buckets, got {len(rg.data)}" + + def test_all_sizes_within_bounds(self): + rg = _build_plain_group() + for r in rg.data: + assert _MIN_SIDE <= r.h <= _MAX_SIDE, f"height {r.h} outside [{_MIN_SIDE}, {_MAX_SIDE}]" + assert _MIN_SIDE <= r.w <= _MAX_SIDE, f"width {r.w} outside [{_MIN_SIDE}, {_MAX_SIDE}]" + + def test_all_sizes_aligned(self): + rg = _build_plain_group() + for r in rg.data: + assert r.h % _ALIGN == 0, f"height {r.h} not {_ALIGN}-aligned" + assert r.w % _ALIGN == 0, f"width {r.w} not {_ALIGN}-aligned" + + # -- ratio ordering ------------------------------------------------------ + + def test_sorted_by_ratio_ascending(self): + rg = _build_plain_group() + ratios = [r.ratio for r in rg.data] + for i in range(1, len(ratios)): + assert ratios[i - 1] <= ratios[i], f"ratio[{i - 1}]={ratios[i - 1]:.4f} > ratio[{i}]={ratios[i]:.4f}" + + def test_ratio_index_matches_numpy_array(self): + rg = _build_plain_group() + for i, r in enumerate(rg.data): + assert math.isclose(rg.ratio[i], r.ratio, rel_tol=1e-9), ( + f"rg.ratio[{i}]={rg.ratio[i]:.6f} != data[{i}].ratio={r.ratio:.6f}" + ) + + # -- key index positions ------------------------------------------------- + + @pytest.mark.parametrize( + "idx,expected_hw", + [ + # First (most portrait) → 1:4 ratio bucket + (0, (512, 2048)), + # Centre square bucket + (16, (1024, 1024)), + # Most landscape bucket → 4:1 ratio + (32, (2048, 512)), + ], + ) + def test_landmark_index(self, idx, expected_hw): + rg = _build_plain_group() + r = rg.data[idx] + assert (r.h, r.w) == expected_hw, f"idx {idx}: expected {expected_hw}, got ({r.h}, {r.w})" + + def test_center_is_square(self): + rg = _build_plain_group() + r = rg.data[16] + assert r.h == r.w == _BASE + + def test_first_and_last_are_extremes(self): + rg = _build_plain_group() + # first → 512×2048 (ratio 0.25) + assert rg.data[0].ratio < 0.3 + # last → 2048×512 (ratio 4.0) + assert rg.data[-1].ratio > 3.9 + + # -- width × height matching --------------------------------------------- + + def test_every_resolution_has_same_area_product(self): + """Symmetry: stepping loop produces the same product set in both + directions, so (h × w) tuples should appear symmetrically.""" + rg = _build_plain_group() + products = {(r.h, r.w) for r in rg.data} + # For every (h, w) there should exist (w, h) or a nearby equivalent + for h, w in products: + if h == w: + continue # square — trivially symmetric + assert (w, h) in products or any(abs(rr.h - w) <= _STEP and abs(rr.w - h) <= _STEP for rr in rg.data), ( + f"no symmetric counterpart for {h}×{w}" + ) + + def test_product_of_wh_divisible_by_256(self): + """Each bucket's area should be a multiple of 256 (16×16), + which is the VAE down-sample grid.""" + rg = _build_plain_group() + for r in rg.data: + assert (r.h * r.w) % 256 == 0, f"{r.h}×{r.w} = {r.h * r.w} not divisible by 256" + + # -- custom base_size / step / align ------------------------------------- + + @pytest.mark.parametrize("base_size", [512, 768, 1280]) + def test_count_scales_with_base_size(self, base_size): + step = base_size // 16 + rg = ResolutionGroup(base_size=base_size, step=step, align=16, extra_resolutions=None) + expected = 1 + 2 * (base_size // step) # centre + 2×phases + assert len(rg.data) == expected, f"base_size={base_size}: expected {expected}, got {len(rg.data)}" + + def test_custom_step(self): + """A smaller step produces more buckets.""" + rg_small = ResolutionGroup(base_size=1024, step=32, align=16, extra_resolutions=None) + rg_default = _build_plain_group() + assert len(rg_small.data) > len(rg_default.data), "smaller step should produce more buckets" + + +# --------------------------------------------------------------------------- +# 2. HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS matching +# --------------------------------------------------------------------------- + +_EXTRA_SQUARES = {"512x512", "640x640", "768x768", "896x896"} +_EXTRA_NON_SQUARES = {"1024x768", "1280x720", "768x1024", "720x1280"} + + +def _build_extra_group() -> ResolutionGroup: + return ResolutionGroup( + base_size=_BASE, + extra_resolutions=[_make_reso(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS], + ) + + +class TestExtraResolutionMatching: + """Verify extra resolutions are correctly integrated into + the ``ResolutionGroup``.""" + + # -- square extras (ratio = 1.0) ----------------------------------------- + + def test_square_extras_appended_to_base_square(self): + """512², 640², 768², 896² all have ratio 1.0 and should be + appended to the (1024, 1024) Resolution via ``append()``.""" + rg = _build_extra_group() + # Find the 1024×1024 entry + square_reso = next(r for r in rg.data if r.h == 1024 and r.w == 1024) + expected_extra = {(512, 512), (640, 640), (768, 768), (896, 896)} + assert square_reso.extra_res == expected_extra, f"square extra_res: {square_reso.extra_res}" + + def test_base_1024_unchanged(self): + """The 1024×1024 entry itself still reports h=1024, w=1024.""" + rg = _build_extra_group() + r = next(r for r in rg.data if r.h == 1024 and r.w == 1024) + assert r.h == 1024 and r.w == 1024 + + def test_squares_not_standalone_entries(self): + """Square extras must NOT appear as top-level ``Resolution`` entries; + they live inside ``extra_res`` only.""" + rg = _build_extra_group() + standalone = {(r.h, r.w) for r in rg.data} + for square in [(512, 512), (640, 640), (768, 768), (896, 896)]: + assert square not in standalone, f"{square} should be in extra_res, not a standalone entry" + + # -- non-square extras (ratio ≠ 1.0) ------------------------------------ + + @pytest.mark.parametrize("reso_str", sorted(_EXTRA_NON_SQUARES)) + def test_non_square_extra_is_standalone_entry(self, reso_str): + """1024×768, 1280×720, 768×1024, 720×1280 are appended as + new ``Resolution`` entries because no training bucket shares + their exact ratio.""" + r = _make_reso(reso_str) + rg = _build_extra_group() + matching = [d for d in rg.data if d.h == r.h and d.w == r.w] + assert len(matching) == 1, f"{reso_str} should be in rg.data exactly once" + assert matching[0].extra_res == set(), f"{reso_str} should have empty extra_res, got {matching[0].extra_res}" + + def test_non_square_extras_empty_extra_res(self): + """Non-square extras have no own extra_res.""" + rg = _build_extra_group() + for r in rg.data: + if (r.h, r.w) in {(1024, 768), (1280, 720), (768, 1024), (720, 1280)}: + assert r.extra_res == set() + + # -- total count --------------------------------------------------------- + + def test_total_count(self): + """33 training + 4 non-square extras = 37.""" + rg = _build_extra_group() + assert len(rg.data) == 37, f"expected 37, got {len(rg.data)}" + + def test_extra_count(self): + """Exactly 8 extra resolutions (4 square + 4 non-square).""" + rg = _build_extra_group() + reso_set = _reso_set(rg) + extra_in_set = reso_set - {(r.h, r.w) for r in _build_plain_group().data} + assert len(extra_in_set) == 8, f"expected 8 extra, got {len(extra_in_set)}: {extra_in_set}" + + # -- index stability ----------------------------------------------------- + + def test_indices_0_32_unchanged(self): + """Indices 0–32 remain the training buckets (sorted by ratio).""" + plain = _build_plain_group() + rg = _build_extra_group() + for i in range(33): + assert (rg.data[i].h, rg.data[i].w) == (plain.data[i].h, plain.data[i].w), ( + f"idx {i}: extra group has ({rg.data[i].h},{rg.data[i].w}), " + f"expected ({plain.data[i].h},{plain.data[i].w})" + ) + + @pytest.mark.parametrize( + "reso_str,expected_idx", + [ + ("1024x768", 33), + ("1280x720", 34), + ("768x1024", 35), + ("720x1280", 36), + ], + ) + def test_non_square_extra_indices(self, reso_str, expected_idx): + """Non-square extras are appended in order at indices 33-36.""" + r = _make_reso(reso_str) + rg = _build_extra_group() + idx = next(i for i, d in enumerate(rg.data) if d.h == r.h and d.w == r.w) + assert idx == expected_idx, f"{reso_str}: expected idx {expected_idx}, got {idx}" + + # -- ratio lookup via get_base_size_and_ratio_index ---------------------- + + @pytest.mark.parametrize( + "h,w,expected_idx", + [ + (1024, 768, 33), + (1280, 720, 34), + (768, 1024, 35), + (720, 1280, 36), + (1024, 1024, 16), + (512, 512, 16), + (640, 640, 16), + (768, 768, 16), + (896, 896, 16), + ], + ) + def test_ratio_index_for_resolutions(self, h, w, expected_idx): + """``get_base_size_and_ratio_index`` must return the correct + ratio index for both training-bucket and extra resolutions.""" + rg = _build_extra_group() + _, idx = rg.get_base_size_and_ratio_index(w, h) + assert idx == expected_idx, f"({w},{h}): expected idx {expected_idx}, got {idx}" + + +# --------------------------------------------------------------------------- +# Resolution.match — area-based matching within a ratio group +# --------------------------------------------------------------------------- + + +class TestResolutionMatch: + """Verify ``Resolution.match`` picks the best-fit alternative from + ``extra_res`` based on area proximity.""" + + def test_no_extra_returns_self(self): + r = _make_reso("1024x768") + assert r.match(800, 600) == (768, 1024) + + @pytest.mark.parametrize( + "target_w,target_h,expected", + [ + # Exact area match in extra_res + (512, 512, (512, 512)), + (640, 640, (640, 640)), + (768, 768, (768, 768)), + (896, 896, (896, 896)), + # Closer to 1024² than to any extra square + (1024, 1024, (1024, 1024)), + (1000, 1000, (1024, 1024)), + # Small square picks 512² + (256, 256, (512, 512)), + (400, 400, (512, 512)), + # 700² → closest to 640² (diff 129600) vs 768² (diff 99856) + # 768*768 - 700*700 = 589824 - 490000 = 99824 + # 640*640 - 700*700 = 409600 - 490000 = -80400 → abs=80400 + # So 640² is closer in area + (700, 700, (640, 640)), + # 800² → closest to 768² (diff 49984) vs 896² (diff 162944) + (800, 800, (768, 768)), + ], + ) + def test_square_extra_match(self, target_w, target_h, expected): + """Build a Resolution like the 1024² entry with square extras, + then verify area-based matching.""" + r = Resolution(1024, 1024) + for s in [(512, 512), (640, 640), (768, 768), (896, 896)]: + r.extra_res.add(s) + result = r.match(target_w, target_h) + assert result == expected, f"match({target_w},{target_h}) = {result}, expected {expected}" + + def test_match_returns_tuple(self): + r = Resolution(1024, 1024) + r.extra_res.add((512, 512)) + result = r.match(512, 512) + assert isinstance(result, tuple) + assert len(result) == 2 + assert all(isinstance(v, int) for v in result) + + def test_match_does_not_mutate_resolution(self): + r = Resolution(1024, 1024) + r.extra_res.add((512, 512)) + extra_before = set(r.extra_res) + r.match(300, 300) + assert r.extra_res == extra_before + + +# --------------------------------------------------------------------------- +# ResolutionGroup.get_target_size — end-to-end resolution selection +# --------------------------------------------------------------------------- + + +class TestGetTargetSize: + """Verify ``get_target_size`` uses ratio-then-area logic correctly.""" + + @pytest.fixture(scope="class") + def rg(self) -> ResolutionGroup: + return _build_extra_group() + + @pytest.mark.parametrize( + "w,h,expected_min_area,expected_max_area", + [ + # Square targets should land on square variants + (1024, 1024, 1024 * 1024, 1024 * 1024), + (512, 512, 512 * 512, 512 * 512), + (640, 640, 640 * 640, 640 * 640), + (768, 768, 768 * 768, 768 * 768), + (896, 896, 896 * 896, 896 * 896), + ], + ) + def test_square_target_returns_square(self, rg, w, h, expected_min_area, expected_max_area): + rw, rh = rg.get_target_size(w, h) + area = rw * rh + assert expected_min_area <= area <= expected_max_area, f"({w},{h}) → ({rw},{rh}) area={area}" + + def test_training_bucket_self_maps(self, rg): + """A training bucket should map to itself (or an extra_res variant + with the same ratio).""" + for r in rg.data: + rw, rh = rg.get_target_size(r.w, r.h) + # Either exact match or the ratio must be preserved + assert abs(rh / rw - r.ratio) < 0.01, f"({r.w},{r.h}) → ({rw},{rh}) ratio drifted from {r.ratio:.4f}" + + def test_wild_ratio_still_finds_bucket(self, rg): + """Even an unusual aspect ratio should map to some valid bucket.""" + for w, h in [(800, 600), (1200, 900), (1600, 400), (400, 1600)]: + rw, rh = rg.get_target_size(w, h) + assert rw > 0 and rh > 0 + assert rw % 16 == 0 and rh % 16 == 0 + + +# --------------------------------------------------------------------------- +# 3. get_cached_resolution_group & ar2diffusion integration +# --------------------------------------------------------------------------- + + +class TestGetCachedResolutionGroup: + """Verify ``get_cached_resolution_group`` returns a correctly + configured ``ResolutionGroup`` and respects the LRU cache.""" + + def test_returns_resolution_group(self): + rg = get_cached_resolution_group(base_size=1024) + assert isinstance(rg, ResolutionGroup) + + def test_base_size_propagated(self): + rg = get_cached_resolution_group(base_size=1024) + assert rg.base_size == 1024 + + def test_has_extra_resolutions(self): + rg = get_cached_resolution_group(base_size=1024) + reso_set = _reso_set(rg) + for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS: + r = _make_reso(s) + assert (r.h, r.w) in reso_set, f"{s} missing from cached group" + + def test_cache_same_base_size_same_object(self): + """LRU cache: same base_size → same object identity.""" + # Clear any prior cache entries + get_cached_resolution_group.cache_clear() + rg1 = get_cached_resolution_group(base_size=1024) + rg2 = get_cached_resolution_group(base_size=1024) + assert rg1 is rg2, "same base_size must return cached object" + + def test_cache_different_base_size_different_object(self): + get_cached_resolution_group.cache_clear() + rg1 = get_cached_resolution_group(base_size=1024) + rg2 = get_cached_resolution_group(base_size=768) + assert rg1 is not rg2, "different base_size must return different objects" + + def test_cache_size_limit(self): + """Cache is bounded at 4 entries (``maxsize=4``).""" + get_cached_resolution_group.cache_clear() + for bs in [512, 768, 1024, 1280, 1536]: + get_cached_resolution_group(base_size=bs) + info = get_cached_resolution_group.cache_info() + assert info.currsize <= 4, f"cache grew to {info.currsize}, max is 4" + # All 5 unique base_size values caused cache misses + assert info.misses >= 5, f"expected >=5 misses for 5 unique base_size values, got {info.misses}" + + def test_resolution_indexing_works(self): + """Indices 0..len(rg)-1 must be accessible.""" + rg = get_cached_resolution_group(base_size=1024) + for i in range(len(rg)): + r = rg[i] + assert isinstance(r, Resolution) + assert r.h > 0 and r.w > 0 + + @pytest.mark.parametrize("base_size", [512, 1024]) + def test_ratio_index_roundtrip_via_group_length(self, base_size): + """Every ratio index in [0, len(rg)) must resolve to a valid + Resolution with height/width within bounds.""" + rg = get_cached_resolution_group(base_size=base_size) + for i in range(len(rg)): + r = rg[i] + assert isinstance(r.height, int) and r.height > 0 + assert isinstance(r.width, int) and r.width > 0 + + +# --------------------------------------------------------------------------- +# ar2diffusion — integration tests +# --------------------------------------------------------------------------- + + +# Token IDs matching the real prompt_utils values +_RATIO_0 = 128044 +_RATIO_32 = 128076 +_RATIO_33 = 130103 +_RATIO_36 = 130106 + + +def _fake_ar_output( + *, + cumulative_token_ids: list[int] | None = None, + text: str = "", +) -> MagicMock: + """Minimal stub for a vLLM ``RequestOutput`` with one completion.""" + choice = MagicMock() + choice.cumulative_token_ids = cumulative_token_ids or [] + choice.text = text + output = MagicMock() + output.outputs = [choice] + return output + + +class TestAr2DiffusionIntegration: + """Verify that ``ar2diffusion`` calls ``get_cached_resolution_group`` + and obtains a correctly configured ``ResolutionGroup`` whose indices map + to the expected training-bucket or extra-resolution sizes.""" + + def _call_ar2diffusion( + self, + token_ids: list[int], + prompt_kwargs: dict[str, Any] | None = None, + ) -> tuple[Any, Any]: + """Invoke ``ar2diffusion`` with a single fake AR output. + + Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can + assert on the resolved height/width. + """ + from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import ( + ar2diffusion, + ) + + pk = dict(prompt_kwargs or {}) + pk.setdefault("prompt", "a beautiful sunset") + pk.setdefault("height", 1024) + pk.setdefault("width", 1024) + pk.setdefault("image_base_size", 1024) + + # Build a prompt-like object that behaves like the real OmniTokensPrompt + class _FakePrompt: + def __init__(self, d): + self._d = d + + def __getitem__(self, key): + return self._d[key] + + def get(self, key, default=None): + return self._d.get(key, default) + + def _asdict(self): + return dict(self._d) + + prompt = _FakePrompt(pk) + ar_output = _fake_ar_output(cumulative_token_ids=token_ids) + result = ar2diffusion( + source_outputs=[ar_output], + prompt=[prompt], + requires_multimodal_data=False, + ) + return result + + # -- ratio token ↔ resolution resolution --------------------------------- + + @pytest.mark.parametrize( + "ratio_idx,expected_hw", + [ + # Training buckets + (0, (512, 2048)), + (16, (1024, 1024)), + (32, (2048, 512)), + # Non-square extras + (33, (1024, 768)), + (34, (1280, 720)), + (35, (768, 1024)), + (36, (720, 1280)), + ], + ) + def test_ratio_index_resolves_to_correct_size(self, ratio_idx, expected_hw): + """When the AR emits ````, the DiT input must carry + the training-bucket width/height for that index.""" + rg = get_cached_resolution_group(base_size=1024) + reso = rg[ratio_idx] + assert (reso.height, reso.width) == expected_hw, ( + f"ratio_idx {ratio_idx}: expected {expected_hw}, got ({reso.height}, {reso.width})" + ) + + @pytest.mark.parametrize("ratio_idx", [0, 16, 32, 33, 34, 35, 36]) + def test_ar2diffusion_extracts_height_width(self, ratio_idx): + """End-to-end: simulate AR output containing ```` + and assert the DiT inputs carry the expected size.""" + rg = get_cached_resolution_group(base_size=1024) + expected_h = rg[ratio_idx].height + expected_w = rg[ratio_idx].width + + token_id = _RATIO_0 + ratio_idx if ratio_idx <= 32 else _RATIO_33 + (ratio_idx - 33) + result = self._call_ar2diffusion( + token_ids=[150000, 150001, token_id], + prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024}, + ) + assert len(result) == 1 + diffusion_input = result[0] + assert diffusion_input["height"] == expected_h, ( + f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}" + ) + assert diffusion_input["width"] == expected_w, ( + f"ratio_idx={ratio_idx}: expected width {expected_w}, got {diffusion_input['width']}" + ) + + def test_no_ratio_token_falls_back_to_prompt_size(self): + """If no ratio token is present, height/width from the prompt + are kept unchanged.""" + result = self._call_ar2diffusion( + token_ids=[150000, 150001], # no ratio tokens + prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024}, + ) + assert len(result) == 1 + diffusion_input = result[0] + assert diffusion_input["height"] == 768 + assert diffusion_input["width"] == 1024 + + def test_out_of_range_ratio_index_falls_back(self): + """If ratio_idx ≥ len(reso_group), ar2diffusion must keep the + prompt size (fallback path).""" + with patch( + "vllm_omni.model_executor.stage_input_processors.hunyuan_image3._extract_ratio_index", + return_value=999, # out of range + ): + result = self._call_ar2diffusion( + token_ids=[150000], + prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024}, + ) + assert len(result) == 1 + assert result[0]["height"] == 512 + assert result[0]["width"] == 512 + + # -- cache behaviour ----------------------------------------------------- + + def test_ar2diffusion_uses_cached_group(self): + """Each invocation of ``ar2diffusion`` calls + ``get_cached_resolution_group``, which returns a shared cached object.""" + get_cached_resolution_group.cache_clear() + rg_before = get_cached_resolution_group(base_size=1024) + + token_id = _RATIO_0 + 16 # = 1024×1024 + self._call_ar2diffusion( + token_ids=[token_id], + prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024}, + ) + + rg_after = get_cached_resolution_group(base_size=1024) + assert rg_before is rg_after, "get_cached_resolution_group should return the same cached object" + + def test_correct_reso_group_used_for_base_size(self): + """Verify ``ar2diffusion`` constructs or retrieves the + ResolutionGroup for the correct ``image_base_size``.""" + get_cached_resolution_group.cache_clear() + + # First request with a non-default base_size + token_id = _RATIO_0 # + self._call_ar2diffusion( + token_ids=[token_id], + prompt_kwargs={"height": 512, "width": 512, "image_base_size": 512}, + ) + + # The cached group for base_size=512 should exist + rg512 = get_cached_resolution_group(base_size=512) + assert rg512.base_size == 512 + + # And base_size=1024 should also be separately accessible + rg1024 = get_cached_resolution_group(base_size=1024) + assert rg1024.base_size == 1024 + assert rg512 is not rg1024 + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- + + +class TestEdgeCases: + def test_resolution_string_parsing(self): + """Resolution accepts both "HxW" strings and (h, w) tuples.""" + r1 = Resolution("512x2048") + assert r1.h == 512 and r1.w == 2048 + + r2 = Resolution(768, 1024) + assert r2.h == 768 and r2.w == 1024 + + r3 = Resolution(1024) # single int → square + assert r3.h == 1024 and r3.w == 1024 + + def test_resolution_extra_res_initially_empty(self): + r = Resolution(1024, 1024) + assert r.extra_res == set() + assert r.match(512, 512) == (1024, 1024) + + def test_resolution_append_deduplicates(self): + """Appending the same resolution twice is a no-op on the set.""" + r = Resolution(1024, 1024) + r.append(Resolution(512, 512)) + r.append(Resolution(512, 512)) + assert len(r.extra_res) == 1 + + def test_all_final_resolutions_are_16_aligned(self): + """Every reachable resolution (including extras) must be + divisible by the VAE down-sample factor of 16.""" + rg = _build_extra_group() + for h, w in _reso_set(rg): + assert h % 16 == 0, f"height {h} not 16-aligned in {h}×{w}" + assert w % 16 == 0, f"width {w} not 16-aligned in {h}×{w}" + + def test_ratio_array_length_matches_data(self): + rg = _build_extra_group() + assert len(rg.ratio) == len(rg.data) + + def test_resolution_repr(self): + r = Resolution(1024, 1024) + rep = repr(r) + assert "1024x1024" in rep + + r2 = Resolution(1024, 1024) + r2.extra_res.add((512, 512)) + rep2 = repr(r2) + assert "512x512" in rep2 diff --git a/tests/diffusion/models/hunyuan_video/__init__.py b/tests/diffusion/models/hunyuan_video/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/tests/diffusion/models/hunyuan_video/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/diffusion/models/hunyuan_video/test_hunyuan_video_quant_config_propagation.py b/tests/diffusion/models/hunyuan_video/test_hunyuan_video_quant_config_propagation.py new file mode 100644 index 000000000..524a9b93b --- /dev/null +++ b/tests/diffusion/models/hunyuan_video/test_hunyuan_video_quant_config_propagation.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for HunyuanVideo-1.5 quant_config propagation through transformer creation. + +Tests cover: +- HunyuanVideo15Pipeline passes quant_config to HunyuanVideo15Transformer3DModel +- HunyuanVideo15I2VPipeline passes quant_config to HunyuanVideo15Transformer3DModel +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 as t2v_module +import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v as i2v_module + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class TestHunyuanVideoQuantConfigPropagation: + """Verify quant_config is propagated to the transformer model in HunyuanVideo-1.5 pipelines.""" + + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_local_device", + return_value="cpu", + ) + @patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.prefetch_subfolders") + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.from_pretrained_with_prefetch", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.Qwen2Tokenizer.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.ByT5Tokenizer.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.AutoConfig.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.T5EncoderModel", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.FlowMatchEulerDiscreteScheduler.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_transformer_config_kwargs", + return_value={}, + ) + def test_t2v_quant_config_passed_through( + self, + mock_get_kwargs, + mock_scheduler, + mock_t5_encoder, + mock_auto_config, + mock_byt5_tok, + mock_qwen2_tok, + mock_from_pretrained, + mock_prefetch, + mock_get_device, + ): + captured_kwargs = {} + + class FakeTransformer: + def __init__(self, **kwargs): + captured_kwargs.update(kwargs) + + with patch.object(t2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer): + fake_qc = MagicMock() + od_config = SimpleNamespace( + model="fake-t2v-model", + dtype=torch.bfloat16, + flow_shift=None, + tf_model_config=SimpleNamespace(), + quantization_config=fake_qc, + enable_diffusion_pipeline_profiler=False, + ) + + t2v_module.HunyuanVideo15Pipeline(od_config=od_config) + + assert captured_kwargs.get("quant_config") is fake_qc + + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_local_device", + return_value="cpu", + ) + @patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.prefetch_subfolders") + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.from_pretrained_with_prefetch", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.Qwen2Tokenizer.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.ByT5Tokenizer.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.AutoConfig.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.T5EncoderModel", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.SiglipImageProcessor.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.FlowMatchEulerDiscreteScheduler.from_pretrained", + return_value=MagicMock(), + ) + @patch( + "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_transformer_config_kwargs", + return_value={}, + ) + def test_i2v_quant_config_passed_through( + self, + mock_get_kwargs, + mock_scheduler, + mock_siglip_processor, + mock_t5_encoder, + mock_auto_config, + mock_byt5_tok, + mock_qwen2_tok, + mock_from_pretrained, + mock_prefetch, + mock_get_device, + ): + captured_kwargs = {} + + class FakeTransformer: + def __init__(self, **kwargs): + captured_kwargs.update(kwargs) + + with patch.object(i2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer): + fake_qc = MagicMock() + od_config = SimpleNamespace( + model="fake-i2v-model", + dtype=torch.bfloat16, + flow_shift=None, + tf_model_config=SimpleNamespace(), + quantization_config=fake_qc, + enable_diffusion_pipeline_profiler=False, + ) + + i2v_module.HunyuanVideo15I2VPipeline(od_config=od_config) + + assert captured_kwargs.get("quant_config") is fake_qc diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index 3930677d3..bf6a75636 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py @@ -72,6 +72,285 @@ class TestPipelineIndependence: assert issubclass(LTX23Pipeline, CFGParallelMixin) + def test_ltx23_pipeline_declares_offload_components(self): + """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + + pipe = object.__new__(LTX23Pipeline) + torch.nn.Module.__init__(pipe) + pipe.transformer = torch.nn.Linear(1, 1) + pipe.text_encoder = torch.nn.Linear(1, 1) + pipe.connectors = torch.nn.Linear(1, 1) + pipe.vae = torch.nn.Linear(1, 1) + pipe.audio_vae = torch.nn.Linear(1, 1) + pipe.vocoder = torch.nn.Linear(1, 1) + + modules = ModuleDiscovery.discover(pipe) + + assert LTX23Pipeline._dit_modules == ["transformer"] + assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"] + assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"] + assert LTX23Pipeline._resident_modules == ["vocoder"] + assert modules.dit_names == ["transformer"] + assert modules.encoder_names == ["text_encoder", "connectors"] + assert modules.resident_names == ["vocoder"] + assert len(modules.vaes) == 2 + + def test_ltx23_pipeline_has_diffusion_pipeline_profiler_mixin(self): + """LTX23Pipeline must support lightweight stage timing.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin + + assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin) + + +class TestLTX23VaeDecodeParallel: + """Test LTX-2.3 video VAE tiled parallel helpers without loading weights.""" + + def test_ltx23_video_vae_is_distributed_tile_only_class(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin + + assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin) + assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split") + + def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + + vae = SimpleNamespace( + spatial_compression_ratio=32, + tile_sample_min_height=512, + tile_sample_min_width=512, + tile_sample_stride_height=448, + tile_sample_stride_width=448, + temporal_compression_ratio=8, + dtype=torch.float32, + ) + + z = torch.zeros(1, 2, 5, 16, 24) + tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z) + + assert grid_spec.grid_shape == (2, 2) + assert grid_spec.split_dims == (3, 4) + assert grid_spec.tile_spec["sample_height"] == 512 + assert grid_spec.tile_spec["sample_width"] == 768 + assert grid_spec.tile_spec["blend_height"] == 64 + assert grid_spec.tile_spec["blend_width"] == 64 + assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512) + assert grid_spec.tile_spec["tile_output_shapes"] == { + 0: (1, 3, 33, 512, 512), + 1: (1, 3, 33, 512, 320), + 2: (1, 3, 33, 64, 512), + 3: (1, 3, 33, 64, 320), + } + assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)] + assert [tuple(task.tensor.shape) for task in tasks] == [ + (1, 2, 5, 16, 16), + (1, 2, 5, 16, 10), + (1, 2, 5, 2, 16), + (1, 2, 5, 2, 10), + ] + assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10] + + def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec + + class FakeVae: + def __init__(self): + self.blend_calls = [] + + def clear_cache(self): + pass + + def blend_v(self, _previous, current, blend_height): + self.blend_calls.append(("v", blend_height)) + return current + + def blend_h(self, _previous, current, blend_width): + self.blend_calls.append(("h", blend_width)) + return current + + fake_vae = FakeVae() + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(2, 2), + tile_spec={ + "sample_height": 10, + "sample_width": 10, + "blend_height": 1, + "blend_width": 2, + "tile_sample_stride_height": 5, + "tile_sample_stride_width": 5, + }, + ) + tiles = { + (0, 0): torch.full((1, 3, 2, 6, 6), 1.0), + (0, 1): torch.full((1, 3, 2, 6, 6), 2.0), + (1, 0): torch.full((1, 3, 2, 6, 6), 3.0), + (1, 1): torch.full((1, 3, 2, 6, 6), 4.0), + } + + merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec) + + assert merged.shape == (1, 3, 2, 10, 10) + assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)] + torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5)) + torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0)) + torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0)) + torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0)) + + def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self): + from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2 + + z = torch.zeros(1, 2, 1, 16, 24) + expected = torch.ones(1, 3, 1, 512, 768) + seen = {} + + class FakeExecutor: + def execute(self, tensor, operator, broadcast_result=True): + seen["tensor"] = tensor + seen["operator"] = operator + seen["broadcast_result"] = broadcast_result + return expected + + vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True) + vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae) + vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae) + vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae) + + output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode( + vae, + z, + temb=torch.tensor(0.5), + return_dict=False, + ) + + assert len(output) == 1 + assert output[0] is expected + assert seen["tensor"] is z + assert seen["broadcast_result"] is False + assert seen["operator"].split.__name__ == "tile_split" + assert seen["operator"].merge.__name__ == "tile_merge" + + def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor + from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + GridSpec, + TileTask, + ) + + z = torch.zeros(1, 1, 1, 1, 1) + tile_output_shapes = { + 0: (1, 1, 1, 2, 2), + 1: (1, 1, 1, 2, 1), + 2: (1, 1, 1, 1, 2), + 3: (1, 1, 1, 1, 1), + } + tasks = [ + TileTask(0, (0, 0), z, workload=4), + TileTask(1, (0, 1), z, workload=2), + TileTask(2, (1, 0), z, workload=2), + TileTask(3, (1, 1), z, workload=1), + ] + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(2, 2), + tile_spec={ + "max_tile_output_shape": (1, 1, 1, 2, 2), + "tile_output_shapes": tile_output_shapes, + }, + output_dtype=torch.float32, + ) + seen = {} + + def exec_tile(task): + return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1)) + + def merge_tiles(coord_tensor_map, passed_grid_spec): + seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()} + assert passed_grid_spec is grid_spec + return torch.stack( + [ + coord_tensor_map[(0, 0)].flatten()[0], + coord_tensor_map[(0, 1)].flatten()[0], + coord_tensor_map[(1, 0)].flatten()[0], + coord_tensor_map[(1, 1)].flatten()[0], + ] + ) + + operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles) + + rank0_executor = object.__new__(LTX2VaeExecutor) + rank0_executor.parallel_size = 2 + rank0_executor.world_size = 2 + rank0_executor.rank = 0 + + def gather_rank0(local_tile_tensor): + assigned = rank0_executor._balance_tasks(tasks, 2) + rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]] + rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta( + rank1_results, + list(local_tile_tensor.shape), + z.device, + torch.float32, + ) + seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape) + return [local_tile_tensor, rank1_tile_tensor] + + def fail_final_sync(*_args, **_kwargs): + raise AssertionError("broadcast_result=False should not sync the final result") + + rank0_executor.gather_tensors = gather_rank0 + rank0_executor._sync_final_result = fail_final_sync + + rank0_result = rank0_executor.execute(z, operator, broadcast_result=False) + + torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0])) + assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2) + assert seen["merged_shapes"] == { + (0, 0): (1, 1, 1, 2, 2), + (0, 1): (1, 1, 1, 2, 1), + (1, 0): (1, 1, 1, 1, 2), + (1, 1): (1, 1, 1, 1, 1), + } + + non_rank0_executor = object.__new__(LTX2VaeExecutor) + non_rank0_executor.parallel_size = 2 + non_rank0_executor.world_size = 2 + non_rank0_executor.rank = 1 + + def gather_rank1(local_tile_tensor): + seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape) + return None + + def fail_non_rank0_merge(*_args, **_kwargs): + raise AssertionError("non-rank0 should not merge gathered tiles") + + non_rank0_executor.gather_tensors = gather_rank1 + non_rank0_executor._sync_final_result = fail_final_sync + + empty_result = non_rank0_executor.execute( + z, + DistributedOperator( + split=lambda _z: (tasks, grid_spec), + exec=exec_tile, + merge=fail_non_rank0_merge, + ), + broadcast_result=False, + ) + + assert tuple(empty_result.shape) == (0,) + assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2) + class TestCFGParallelHelpers: """Test LTX-2.3 CFG helper math without loading model weights.""" diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py new file mode 100644 index 000000000..e71253297 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from torch import nn + +from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm + + +def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer(): + norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False) + + assert "weight" not in dict(norm.named_parameters()) + assert "weight" in dict(norm.named_buffers()) + assert "weight" not in norm.state_dict() + + +def test_ltx_rms_norm_affine_weight_remains_parameter(): + norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=True) + + assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter) + assert "weight" not in dict(norm.named_buffers()) + assert "weight" in norm.state_dict() diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index 502a923ba..fa26ae276 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -41,7 +41,7 @@ def mock_od_config(mocker: MockerFixture): config.cache_backend = None config.cache_config = None config.model = "test-model" - config.diffusion_load_format = None + config.diffusion_load_format = "default" config.dtype = "float32" config.max_cpu_loras = 0 config.lora_path = None diff --git a/tests/distributed/omni_connectors/test_adapter_and_flow.py b/tests/distributed/omni_connectors/test_adapter_and_flow.py index c2be5003c..548d16eb6 100644 --- a/tests/distributed/omni_connectors/test_adapter_and_flow.py +++ b/tests/distributed/omni_connectors/test_adapter_and_flow.py @@ -145,7 +145,7 @@ def test_shm_connector_flow(mocker: MockerFixture): Using real SharedMemoryConnector (inline mode for simplicity). """ # 1. Setup Connector - config = {"shm_threshold_bytes": 1024} # Large threshold to use inline + config = {"shm_threshold_bytes": 1024, "inline_small_payloads": True} connector = SharedMemoryConnector(config) connectors_map = {("0", "1"): connector} diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index 4d0661be8..1518e98b1 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/tests/distributed/omni_connectors/test_basic_connectors.py @@ -75,18 +75,21 @@ def test_create_unknown_connector(): @pytest.fixture def shm_connector(): - config = {"shm_threshold_bytes": 100} # Small threshold for testing + config = {"shm_threshold_bytes": 100, "inline_small_payloads": True} return SharedMemoryConnector(config) def test_put_get_inline(shm_connector): - """Test transfer for small data (currently always uses SHM).""" + """Test inline transfer for small data.""" data = {"small": "data"} success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data) assert success is True - assert "shm" in metadata + assert "inline_bytes" in metadata + assert "shm" not in metadata assert "size" in metadata + assert shm_connector._metrics["inline_writes"] == 1 + assert shm_connector._metrics["shm_writes"] == 0 # Retrieve retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index f8eb6af8b..ab671ec66 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -387,14 +387,15 @@ def test_process_and_restore_queues(build_adapter): running_req = _req("r1", RequestStatus.RUNNING) waiting_queue = DummyWaitingQueue([waiting_req]) running_queue = [running_req] + scheduler_requests = {waiting_req.request_id: waiting_req, running_req.request_id: running_req} - adapter.process_pending_chunks(waiting_queue, running_queue) + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK assert running_req.status == RequestStatus.WAITING_FOR_CHUNK assert waiting_queue == [] assert running_queue == [] - adapter.restore_queues(waiting_queue, running_queue) + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests) assert waiting_queue == [waiting_req] assert running_queue == [running_req] assert adapter.waiting_for_chunk_waiting_requests == deque() @@ -897,6 +898,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture): sched = OmniARScheduler.__new__(OmniARScheduler) sched.chunk_transfer_adapter = adapter sched.requests = {} + sched.running = [] + sched.waiting = [] with patch.object(VLLMScheduler, "finish_requests", _super_finish): OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED) @@ -1170,3 +1173,164 @@ def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture): assert eco.outputs[0].finish_reason == "stop" assert eco.outputs[0].kv_transfer_params is None assert scheduler._pending_finish_reqs == [] + + +# --------------------------------------------------------------- +# Zombie purge tests (regression for vllm-project/vllm-omni#3736) +# --------------------------------------------------------------- + + +def test_process_pending_chunks_purges_zombies_in_running_deque( + build_adapter, +): + """A request aborted while parked in ``waiting_for_chunk_running_requests`` + is no longer tracked by ``scheduler.requests`` once + ``Scheduler._free_request`` runs. ``process_pending_chunks`` must drop + that zombie *before* ``restore_queues`` would re-inject it onto the + running queue. + + Regression for https://github.com/vllm-project/vllm-omni/issues/3736 + (engine-core ``KeyError`` on aborts under chunk-transfer pressure). + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + + live_req = _req("live-1", RequestStatus.WAITING_FOR_CHUNK) + zombie_req = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK) + adapter.waiting_for_chunk_running_requests.append(live_req) + adapter.waiting_for_chunk_running_requests.append(zombie_req) + # Mirror state the adapter would carry for an in-flight request so we can + # later assert ``cleanup_receiver`` actually fired against the zombie. + adapter.requests_origin_status[zombie_req.request_id] = RequestStatus.RUNNING + adapter.requests_origin_status[live_req.request_id] = RequestStatus.RUNNING + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + # Simulate the post-abort scheduler: only ``live-1`` is still tracked. + scheduler_requests = {live_req.request_id: live_req} + + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + + # 1. Zombie was popped from the deque. + assert zombie_req not in adapter.waiting_for_chunk_running_requests + # 2. Live request is still in the deque. + assert live_req in adapter.waiting_for_chunk_running_requests + # 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping + # and registers the id as cancelled so a late load/poll is dropped too). + assert zombie_req.request_id not in adapter.requests_origin_status + assert zombie_req.request_id in adapter._cancelled_load_reqs + # 4. Live request's bookkeeping is untouched. + assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING + assert live_req.request_id not in adapter._cancelled_load_reqs + + # 5. ``restore_queues`` (which the scheduler runs in its ``finally`` + # clause) now only re-injects the live request -- the zombie is + # gone, so the worker's ``_update_states`` cannot crash on it. + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + assert running_queue == [live_req] + + +def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter): + """Zombie purge applies symmetrically to the waiting-side deque. + + Regression for https://github.com/vllm-project/vllm-omni/issues/3736 -- + aborted requests that landed in ``waiting_for_chunk_waiting_requests`` + must also be removed before ``restore_queues`` re-injects them, since + ``restore_queues`` uses ``waiting_queue.add_request`` rather than + ``running_queue.extend`` for that path. + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + + live_req = _req("live-w", RequestStatus.WAITING_FOR_CHUNK) + zombie_req = _req("zombie-w", RequestStatus.WAITING_FOR_CHUNK) + adapter.waiting_for_chunk_waiting_requests.append(live_req) + adapter.waiting_for_chunk_waiting_requests.append(zombie_req) + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + scheduler_requests = {live_req.request_id: live_req} + + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + + assert live_req in adapter.waiting_for_chunk_waiting_requests + assert zombie_req not in adapter.waiting_for_chunk_waiting_requests + assert zombie_req.request_id in adapter._cancelled_load_reqs + + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + assert waiting_queue == [live_req] + + +def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter): + """Interleaved live/zombie ordering -- pin the ``popleft`` + append-live + in-place filter at chunk_transfer_adapter._purge_untracked_chunk_requests: + survivor order must match insertion order, and ``cleanup_receiver`` must + run exactly once per zombie. + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + + live1 = _req("live-1", RequestStatus.WAITING_FOR_CHUNK) + zombie1 = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK) + live2 = _req("live-2", RequestStatus.WAITING_FOR_CHUNK) + zombie2 = _req("zombie-2", RequestStatus.WAITING_FOR_CHUNK) + for req in (live1, zombie1, live2, zombie2): + adapter.waiting_for_chunk_running_requests.append(req) + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + scheduler_requests = {live1.request_id: live1, live2.request_id: live2} + + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + + assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2] + assert zombie1.request_id in adapter._cancelled_load_reqs + assert zombie2.request_id in adapter._cancelled_load_reqs + + +def test_restore_queues_purges_late_aborts_after_process_pending_chunks( + build_adapter, +): + """Race window guard: an abort can fire between ``process_pending_chunks`` + and the scheduler's ``finally``-clause ``restore_queues`` call. The + second purge inside ``restore_queues`` must drop the now-untracked + request so it does not get re-injected onto ``running_queue``. + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + + req = _req("late-abort", RequestStatus.WAITING_FOR_CHUNK) + adapter.waiting_for_chunk_running_requests.append(req) + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + + # Tick N: process_pending_chunks runs while req is still tracked. + scheduler_requests: dict = {req.request_id: req} + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + assert req in adapter.waiting_for_chunk_running_requests + + # Mid-tick: abort fires and the scheduler's free path deletes the entry. + del scheduler_requests[req.request_id] + + # finally: restore_queues sees the now-untracked req and must drop it + # instead of blindly extending it onto running_queue. + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + assert running_queue == [] + assert req.request_id in adapter._cancelled_load_reqs + + +def test_purge_is_noop_on_empty_deques(build_adapter): + """Empty deques short-circuit -- guards against any accidental + ``IndexError`` from ``popleft`` on an empty deque if a future caller + shadows the empty check. + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + assert len(adapter.waiting_for_chunk_waiting_requests) == 0 + assert len(adapter.waiting_for_chunk_running_requests) == 0 + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests={}) + assert len(adapter.waiting_for_chunk_waiting_requests) == 0 + assert len(adapter.waiting_for_chunk_running_requests) == 0 + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={}) + assert running_queue == [] + assert waiting_queue == [] diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index e702318e3..73928aaf6 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/tests/distributed/omni_connectors/test_shm_connector.py @@ -13,7 +13,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @pytest.fixture() def connector(): - c = SharedMemoryConnector({"shm_threshold_bytes": 64}) + c = SharedMemoryConnector({"shm_threshold_bytes": 0}) yield c c.close() diff --git a/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py b/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py index 5e872e73a..8ce6f1988 100644 --- a/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py +++ b/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py @@ -340,7 +340,7 @@ def build_arg_parser() -> argparse.ArgumentParser: p.add_argument( "--max-seed-tts-mean-wer", type=float, - default=0.25, + default=0.35, help="If set, fail when seed_tts_content_error_mean is strictly above this value.", ) p.add_argument( diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py index 26b39ea78..ad4dd5726 100644 --- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py +++ b/tests/e2e/accuracy/test_diffusers_backend_similarity.py @@ -218,7 +218,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I torch_dtype=torch.bfloat16, trust_remote_code=True, ).to("cuda") - + pipe.transformer.set_attention_backend("_flash_3_hub") _diffusers_dummy_run(pipe) generator = torch.Generator(device="cuda").manual_seed(SEED) diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index 7ed3c85a0..620254cf9 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -29,7 +29,46 @@ from tests.e2e.accuracy.helpers import ( ) from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner, OmniServer -from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids +from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + build_prompt_tokens, + resolve_stop_token_ids, +) + +# CoT structural markers that MUST appear in AR output text. +# The AR trajectory under think_recaption mode is: +# ... ... ... +# ... +# Test deploy config sets include_stop_str_in_output=True and +# skip_special_tokens=False, so ALL opening and closing markers +# should be present in the detokenized text. +_THINK_CLOSE = "" # avoid Python slash-t-as-TAB pitfall +_COT_REQUIRED_MARKERS = ( + "", + "", + "", + "", + "", + "", + _THINK_CLOSE, +) + + +def _assert_cot_structural_markers(cot_text: str, label: str) -> None: + """Assert that the AR CoT output contains all required structural tags. + + Guards against regressions where detokenizer stop-token stripping + or truncation logic loses CoT content. Under the test deploy config + (include_stop_str_in_output=True, skip_special_tokens=False), + ALL markers should be present in the detokenized text. + """ + missing = [m for m in _COT_REQUIRED_MARKERS if m not in cot_text] + assert not missing, ( + f"[{label}] CoT text missing required structural markers: {missing}. " + f"This typically indicates that stop-token stripping or " + f"_truncate_at_cot_end truncation removed CoT content. " + f"CoT text length={len(cot_text)}, first 100 chars: {repr(cot_text[:100])}" + ) + os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" @@ -129,6 +168,7 @@ _DEPLOY_CONFIG = { "stop_token_ids": [128025], "detokenize": True, "skip_special_tokens": False, + "include_stop_str_in_output": True, }, }, { @@ -400,6 +440,8 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_ ] print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid")) + _assert_cot_structural_markers(online_cot, "ONLINE") + assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], ( f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {THRESHOLDS['cot_semantic_sim']}" ) @@ -508,6 +550,8 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_ print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid")) + _assert_cot_structural_markers(omni_cot, "OFFLINE") + assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], ( f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}" ) diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index c9a09e6de..2b513ecda 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -106,6 +106,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: trust_remote_code=True, local_files_only=_local_files_only(model), ).to("cuda") + pipe.transformer.set_attention_backend("_flash_3_hub") generator = torch.Generator(device="cuda").manual_seed(SEED) result = pipe( # pyright: ignore[reportCallIssue] prompt=PROMPT, @@ -167,6 +168,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im trust_remote_code=True, local_files_only=_local_files_only(model), ).to("cuda") + pipe.transformer.set_attention_backend("_flash_3_hub") generator = torch.Generator(device="cuda").manual_seed(SEED_2512) result = pipe( # pyright: ignore[reportCallIssue] prompt=PROMPT_2512, diff --git a/tests/e2e/accuracy/test_videogen_modelopt_quant.py b/tests/e2e/accuracy/test_videogen_modelopt_quant.py new file mode 100644 index 000000000..e91c15175 --- /dev/null +++ b/tests/e2e/accuracy/test_videogen_modelopt_quant.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end ModelOpt FP8 quality accuracy for video-gen DiTs. + +Modeled on ``test_hunyuan_image3.py``'s ``test_quantized_dit_matches_bf16_accuracy``: +for each calibrated FP8 checkpoint the test generates one video with BF16 and one +with the FP8 DiT under the same seed, scores prompt-faithfulness with CLIP on the +middle frame, then asserts both ``CLIP >= absolute_floor`` and +``CLIP_drop <= drop_threshold``. The test is gated by an opt-in env var (model +paths are local) so it never runs in generic CI. + +To run:: + + VIDEOGEN_RUN_QUANT_ACCURACY=1 \\ + WAN22_A14B_BF16_MODEL=Wan-AI/Wan2.2-T2V-A14B-Diffusers \\ + WAN22_A14B_FP8_MODEL=/path/to/wan22-a14b-modelopt-fp8 \\ + HV15_BF16_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\ + HV15_FP8_MODEL=/path/to/hv15-modelopt-fp8 \\ + pytest -s -v tests/e2e/accuracy/test_videogen_modelopt_quant.py +""" + +from __future__ import annotations + +import gc +import json +import os +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch +from PIL import Image + +from tests.e2e.accuracy.helpers import CLIPScorer, model_output_dir + +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion] + +# --------------------------------------------------------------------------- +# Tunables (override via env to suit local hardware / calibration budget) +# --------------------------------------------------------------------------- +SEED = 42 +HEIGHT = int(os.environ.get("VIDEOGEN_QUANT_HEIGHT", "480")) +WIDTH = int(os.environ.get("VIDEOGEN_QUANT_WIDTH", "832")) +NUM_FRAMES = int(os.environ.get("VIDEOGEN_QUANT_NUM_FRAMES", "17")) +NUM_INFERENCE_STEPS = int(os.environ.get("VIDEOGEN_QUANT_STEPS", "20")) + +WAN_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic." +HV_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic." + +# CLIP gates: absolute floor + max drop from the BF16 reference. +CLIP_ABSOLUTE_FLOOR = float(os.environ.get("VIDEOGEN_QUANT_CLIP_FLOOR", "20.0")) +CLIP_DROP_FP8 = float(os.environ.get("VIDEOGEN_QUANT_CLIP_DROP_FP8", "7.0")) + +RUN_ENV = "VIDEOGEN_RUN_QUANT_ACCURACY" +_TRUE = {"1", "true", "yes", "on"} + + +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class _Spec: + """A model under test (BF16 + its FP8 variant).""" + + family: str + bf16_env: str + fp8_env: str + prompt: str + guidance_scale: float + guidance_scale_2: float | None = None + boundary_ratio: float | None = None + + +WAN_SPEC = _Spec( + family="wan22_a14b_t2v", + bf16_env="WAN22_A14B_BF16_MODEL", + fp8_env="WAN22_A14B_FP8_MODEL", + prompt=WAN_PROMPT, + guidance_scale=3.5, + guidance_scale_2=4.0, + boundary_ratio=0.875, +) + +HV_SPEC = _Spec( + family="hunyuanvideo15_t2v", + bf16_env="HV15_BF16_MODEL", + fp8_env="HV15_FP8_MODEL", + prompt=HV_PROMPT, + guidance_scale=6.0, +) + + +@dataclass(frozen=True) +class _QuantCase: + name: str + spec: _Spec + clip_drop: float + + +def _cases() -> list[pytest.ParameterSet]: + cases = [ + _QuantCase("wan_fp8", WAN_SPEC, CLIP_DROP_FP8), + _QuantCase("hv_fp8", HV_SPEC, CLIP_DROP_FP8), + ] + params: list[pytest.ParameterSet] = [] + opted_in = os.environ.get(RUN_ENV, "").lower() in _TRUE + for case in cases: + marks: list[Any] = [] + if not opted_in: + marks.append(pytest.mark.skip(reason=f"Set {RUN_ENV}=1 to run video-gen quant accuracy.")) + if not os.environ.get(case.spec.bf16_env): + marks.append(pytest.mark.skip(reason=f"Set {case.spec.bf16_env} to run {case.name}.")) + if not os.environ.get(case.spec.fp8_env): + marks.append(pytest.mark.skip(reason=f"Set {case.spec.fp8_env} to run {case.name}.")) + params.append(pytest.param(case, id=case.name, marks=marks)) + return params + + +# --------------------------------------------------------------------------- +def _middle_frame(frames: list[Image.Image] | list[np.ndarray] | np.ndarray) -> Image.Image: + if isinstance(frames, np.ndarray): + idx = frames.shape[0] // 2 + arr = frames[idx] + if arr.dtype != np.uint8: + arr = (arr * 255).clip(0, 255).astype(np.uint8) + return Image.fromarray(arr).convert("RGB") + if not frames: + raise AssertionError("Empty video output") + mid = frames[len(frames) // 2] + if isinstance(mid, np.ndarray): + if mid.dtype != np.uint8: + mid = (mid * 255).clip(0, 255).astype(np.uint8) + return Image.fromarray(mid).convert("RGB") + return mid.convert("RGB") + + +def _extract_frames(result) -> list[Image.Image]: + """Pull a frame list out of an Omni text-to-video result.""" + from vllm_omni.entrypoints.omni import OmniRequestOutput + + if isinstance(result, list): + result = result[0] if result else None + if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None: + result = result.request_output + images = getattr(result, "images", None) + if not images: + raise AssertionError("Pipeline output had no images") + if isinstance(images, list) and images and isinstance(images[0], tuple): + # (frames, fps) + images = images[0][0] + return list(images) + + +def _generate_video( + *, + model: str, + spec: _Spec, + quantization: str | None, + output_path: Path, +) -> tuple[Image.Image, float]: + """Generate one video and return its middle frame + elapsed wall time.""" + from vllm_omni.diffusion.data import DiffusionParallelConfig + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.platforms import current_omni_platform + + kwargs: dict[str, Any] = dict( + model=model, + parallel_config=DiffusionParallelConfig(), + enforce_eager=True, + vae_use_tiling=True, + ) + if quantization: + kwargs["quantization"] = quantization + if spec.boundary_ratio is not None: + kwargs["boundary_ratio"] = spec.boundary_ratio + omni = Omni(**kwargs) + + try: + gen = torch.Generator(device=current_omni_platform.device_type).manual_seed(SEED) + sampling = dict( + height=HEIGHT, + width=WIDTH, + generator=gen, + guidance_scale=spec.guidance_scale, + num_inference_steps=NUM_INFERENCE_STEPS, + num_frames=NUM_FRAMES, + seed=SEED, + ) + if spec.guidance_scale_2 is not None: + sampling["guidance_scale_2"] = spec.guidance_scale_2 + + import time as _time + + t0 = _time.perf_counter() + out = omni.generate({"prompt": spec.prompt}, OmniDiffusionSamplingParams(**sampling)) + elapsed = _time.perf_counter() - t0 + + frames = _extract_frames(out) + mid = _middle_frame(frames) + mid.save(output_path) + return mid, elapsed + finally: + del omni + gc.collect() + if torch.cuda.is_available(): + torch.accelerator.empty_cache() + + +@pytest.mark.parametrize("case", _cases()) +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.accelerator.device_count() < 1, + reason="Needs at least 1 GPU.", +) +def test_quantized_videogen_dit_matches_bf16_clip(case: _QuantCase, accuracy_artifact_root: Path) -> None: + """FP8 DiT should preserve prompt-aligned video quality vs BF16.""" + out_dir = model_output_dir(accuracy_artifact_root, case.spec.family + "-quant") + + with tempfile.TemporaryDirectory(): + bf16_model = os.environ[case.spec.bf16_env] + quant_model = os.environ[case.spec.fp8_env] + + bf16_frame, bf16_time = _generate_video( + model=bf16_model, + spec=case.spec, + quantization=None, + output_path=out_dir / "bf16.png", + ) + quant_frame, quant_time = _generate_video( + model=quant_model, + spec=case.spec, + quantization="fp8", + output_path=out_dir / f"{case.name}.png", + ) + + clip = CLIPScorer() + bf16_clip = clip.score(bf16_frame, case.spec.prompt) + quant_clip = clip.score(quant_frame, case.spec.prompt) + clip_drop = bf16_clip - quant_clip + speedup = (bf16_time / quant_time) if quant_time else float("nan") + + metrics = { + "case": case.name, + "family": case.spec.family, + "quantization": "fp8", + "bf16_model": bf16_model, + "quant_model": quant_model, + "prompt": case.spec.prompt, + "height": HEIGHT, + "width": WIDTH, + "num_frames": NUM_FRAMES, + "num_inference_steps": NUM_INFERENCE_STEPS, + "bf16_elapsed_s": bf16_time, + "quant_elapsed_s": quant_time, + "speedup_vs_bf16": speedup, + "bf16_clip_score": bf16_clip, + "quant_clip_score": quant_clip, + "clip_score_drop": clip_drop, + } + metrics_path = out_dir / f"{case.name}_metrics.json" + metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8") + + print(f"\n[{case.name}] BF16 {bf16_time:.2f}s | FP8 {quant_time:.2f}s (x{speedup:.2f})") + print(f" CLIP BF16={bf16_clip:.2f} FP8={quant_clip:.2f} drop={clip_drop:.2f}") + print(f" metrics={metrics_path}") + + assert quant_clip >= CLIP_ABSOLUTE_FLOOR, ( + f"{case.name} CLIP below floor: got {quant_clip:.2f}, expected >= {CLIP_ABSOLUTE_FLOOR:.2f}" + ) + assert clip_drop <= case.clip_drop, ( + f"{case.name} CLIP drop too large: got {clip_drop:.2f}, expected <= {case.clip_drop:.2f}" + ) diff --git a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py index a81cfe6d4..19cf6fa4c 100644 --- a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py +++ b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py @@ -51,7 +51,7 @@ test_params = [(m, c) for m in models for c in [get_eager_config()]] @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_text_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -67,7 +67,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_image_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -84,7 +84,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_audio_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -103,7 +103,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_video_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -120,7 +120,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -140,7 +140,7 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.advanced_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_text_to_audio(omni_runner, omni_runner_handler) -> None: """ @@ -156,7 +156,7 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None: @pytest.mark.advanced_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_image_to_audio(omni_runner, omni_runner_handler) -> None: """ diff --git a/tests/e2e/offline_inference/test_ming_tts.py b/tests/e2e/offline_inference/test_ming_tts.py new file mode 100644 index 000000000..ca76c9afe --- /dev/null +++ b/tests/e2e/offline_inference/test_ming_tts.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""End-to-end offline inference tests for Ming-omni-tts.""" + +import asyncio +import uuid + +import numpy as np +import pytest +import torch +from transformers import AutoTokenizer +from vllm import SamplingParams +from vllm.inputs import tokens_input + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni import AsyncOmni +from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + KEY_MAX_DECODE_STEPS, + SAMPLE_RATE, + TEXT_EOS_TOKEN_ID, +) +from vllm_omni.model_executor.models.ming_tts.prompt_assembly import DEFAULT_PROMPT, build_ming_dense_prompt + +MODEL = "inclusionAI/Ming-omni-tts-0.5B" +DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml") +TEST_TEXT = "我会一直在这里陪着你，直到你慢慢地沉入那个最温柔的梦里。" +TEST_INSTRUCTION = "轻柔的ASMR耳语，慢速，贴近麦克风" +MIN_AUDIO_SAMPLES = 1000 + + +@pytest.fixture(scope="module") +def ming_tokenizer(): + return AutoTokenizer.from_pretrained(MODEL, trust_remote_code=False) + + +@pytest.fixture +def ming_engine(): + with OmniRunner( + MODEL, + deploy_config=DEPLOY_CONFIG, + stage_init_timeout=300, + enforce_eager=True, + ) as runner: + yield runner.omni + + +@pytest.fixture +def async_omni_engine(): + engine = AsyncOmni( + model=MODEL, + deploy_config=DEPLOY_CONFIG, + stage_init_timeout=300, + enforce_eager=True, + ) + yield engine + engine.shutdown() + + +def _build_prompt( + tokenizer, + *, + text: str = TEST_TEXT, + instruction=TEST_INSTRUCTION, + use_zero_spk_emb: bool = True, +) -> dict: + prompt_dict = build_ming_dense_prompt( + tokenizer, + prompt=DEFAULT_PROMPT, + text=text, + instruction=instruction, + runtime_controls={KEY_MAX_DECODE_STEPS: 200}, + use_zero_spk_emb=use_zero_spk_emb, + ) + prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"]) + prompt["prompt"] = prompt_dict["prompt"] + prompt["text"] = prompt_dict["text"] + prompt["additional_information"] = prompt_dict["additional_information"] + return prompt + + +def _sampling_params_list() -> list[SamplingParams]: + return [ + SamplingParams( + temperature=0.0, + max_tokens=201, + stop_token_ids=[int(TEXT_EOS_TOKEN_ID)], + ), + SamplingParams(temperature=0.0, max_tokens=1), + ] + + +def _flatten_audio(audio) -> torch.Tensor: + if isinstance(audio, list): + parts = [torch.as_tensor(item, dtype=torch.float32).reshape(-1).cpu() for item in audio] + parts = [item for item in parts if item.numel() > 0] + if not parts: + return torch.zeros((0,), dtype=torch.float32) + return torch.cat(parts, dim=0) + return torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu() + + +def _extract_audio(multimodal_output: dict) -> torch.Tensor: + assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}" + audio = multimodal_output.get("audio", multimodal_output.get("model_outputs")) + assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}" + waveform = _flatten_audio(audio) + if waveform.numel() == 0: + raise RuntimeError("Generated audio waveform is empty") + return waveform + + +def _extract_sample_rate(multimodal_output: dict) -> int: + sample_rate = multimodal_output.get("sr") + if sample_rate is None: + raise RuntimeError("Expected multimodal_output['sr']") + if isinstance(sample_rate, list): + sample_rate = sample_rate[-1] + if hasattr(sample_rate, "item"): + sample_rate = sample_rate.item() + return int(sample_rate) + + +def _extract_final_audio_outputs(outputs): + final_outputs = [] + for item in outputs: + if getattr(item, "final_output_type", None) == "audio": + final_outputs.append(item) + continue + request_output = getattr(item, "request_output", None) + if request_output is None: + continue + multimodal_output = getattr(request_output, "multimodal_output", None) + if isinstance(multimodal_output, dict): + final_outputs.append(item) + continue + completions = getattr(request_output, "outputs", None) or [] + if any(isinstance(getattr(completion, "multimodal_output", None), dict) for completion in completions): + final_outputs.append(item) + return final_outputs + + +def _extract_multimodal_output(output) -> dict: + multimodal_output = getattr(output, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + + request_output = getattr(output, "request_output", None) + if request_output is not None: + multimodal_output = getattr(request_output, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + completions = getattr(request_output, "outputs", None) or [] + for completion in completions: + multimodal_output = getattr(completion, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + + raise AssertionError("No multimodal audio output found in Ming generate results") + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ming_tts_offline_basic(ming_engine, ming_tokenizer) -> None: + """Test blocking Ming generation through Omni.""" + outputs = ming_engine.generate( + prompts=[_build_prompt(ming_tokenizer)], + sampling_params_list=_sampling_params_list(), + py_generator=False, + ) + final_outputs = _extract_final_audio_outputs(outputs) + assert len(final_outputs) == 1, f"Expected one final audio output, got {len(final_outputs)}" + multimodal_output = _extract_multimodal_output(final_outputs[0]) + waveform = _extract_audio(multimodal_output) + sample_rate = _extract_sample_rate(multimodal_output) + assert waveform.ndim == 1 + assert waveform.shape[0] == waveform.numel() + assert waveform.numel() > MIN_AUDIO_SAMPLES + assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent" + assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}" + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ming_tts_speaker_conditioning_differs(ming_engine, ming_tokenizer) -> None: + """Test that different Ming speaker controls produce different waveform outputs.""" + style_outputs = ming_engine.generate( + prompts=[_build_prompt(ming_tokenizer)], + sampling_params_list=_sampling_params_list(), + py_generator=False, + ) + ip_outputs = ming_engine.generate( + prompts=[_build_prompt(ming_tokenizer, text=TEST_TEXT, instruction={"IP": "灵小甄"}, use_zero_spk_emb=True)], + sampling_params_list=_sampling_params_list(), + py_generator=False, + ) + + style_final_outputs = _extract_final_audio_outputs(style_outputs) + ip_final_outputs = _extract_final_audio_outputs(ip_outputs) + assert len(style_final_outputs) == 1, "No style audio output produced" + assert len(ip_final_outputs) == 1, "No IP audio output produced" + + style_waveform = _extract_audio(_extract_multimodal_output(style_final_outputs[0])) + ip_waveform = _extract_audio(_extract_multimodal_output(ip_final_outputs[0])) + assert style_waveform.numel() > MIN_AUDIO_SAMPLES + assert ip_waveform.numel() > MIN_AUDIO_SAMPLES + assert np.max(np.abs(style_waveform.numpy())) > 0.01, "Style audio appears silent" + assert np.max(np.abs(ip_waveform.numpy())) > 0.01, "IP audio appears silent" + + overlap = min(int(style_waveform.numel()), int(ip_waveform.numel())) + mean_abs_diff = torch.mean(torch.abs(style_waveform[:overlap] - ip_waveform[:overlap])).item() + assert style_waveform.shape != ip_waveform.shape or mean_abs_diff > 1e-4, ( + "Speaker-conditioned outputs should differ, but style and IP waveforms were effectively identical" + ) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ming_tts_multiple_prompts_queued(ming_engine, ming_tokenizer) -> None: + """Regression: supported max_num_seqs=1 config must still drain queued prompts.""" + prompts = [ + _build_prompt( + ming_tokenizer, + text="第一条语音用于验证 Ming dense 队列中的长请求可以完成。", + instruction="平静自然的旁白，语速中等", + ), + _build_prompt( + ming_tokenizer, + text="第二条短语音也必须生成完成。", + instruction="清晰明亮的女声", + ), + ] + outputs = ming_engine.generate( + prompts=prompts, + sampling_params_list=_sampling_params_list(), + py_generator=False, + ) + final_outputs = _extract_final_audio_outputs(outputs) + assert len(final_outputs) == len(prompts), f"Expected {len(prompts)} audio outputs, got {len(final_outputs)}" + for i, output in enumerate(final_outputs): + waveform = _extract_audio(_extract_multimodal_output(output)) + duration_s = waveform.numel() / SAMPLE_RATE + assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s" + assert np.max(np.abs(waveform.numpy())) > 0.01, f"Request {i} audio appears silent" + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ming_tts_offline_streaming(async_omni_engine, ming_tokenizer) -> None: + """Test async_chunk streaming Ming generation through AsyncOmni.""" + + async def _run() -> None: + all_audio_chunks = [] + accumulated_samples = 0 + chunk_idx = 0 + sample_rate = None + async for stage_output in async_omni_engine.generate( + prompt=_build_prompt(ming_tokenizer), + request_id=str(uuid.uuid4()), + sampling_params_list=_sampling_params_list(), + ): + multimodal_output = stage_output.multimodal_output or {} + audio = multimodal_output.get("audio") + if "sr" in multimodal_output: + sample_rate = _extract_sample_rate(multimodal_output) + if audio is None: + continue + finished = stage_output.finished + if isinstance(audio, torch.Tensor): + if finished: + audio_chunk = audio[accumulated_samples:].float().detach().cpu() + else: + audio_chunk = audio.float().detach().cpu() + elif isinstance(audio, list): + audio_chunk = _flatten_audio(audio) + else: + audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu() + accumulated_samples += int(audio_chunk.numel()) + chunk_idx += 1 + if audio_chunk.numel() > 0: + all_audio_chunks.append(audio_chunk) + assert all_audio_chunks, "No streaming audio chunks received" + waveform = torch.cat(all_audio_chunks, dim=0) + assert waveform.numel() > MIN_AUDIO_SAMPLES + assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent" + assert sample_rate is not None, "Streaming path did not return a sample rate" + assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}" + + asyncio.run(_run()) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py new file mode 100644 index 000000000..1f9769f81 --- /dev/null +++ b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Expanded end-to-end test for SenseNova-U1 in offline mode. + +This test file primarily covers end-to-end tests for LoRA support. + +Validates that LoRA adapters are correctly loaded, applied with controllable +scale, and cleanly deactivated. Uses a synthetic rank-4 adapter targeting the +first decoder layer's gen-path (``_mot_gen``) output projection, so the +perturbation is exercised during the denoising loop. + +Assertions: + (a) LoRA at scale=1.0 visibly changes the output (diff > 0.5) + (b) scale=2.0 produces a different delta than scale=1.0 (scale sensitivity) + (c) The delta is bounded (not corrupted) + (d) Deactivating LoRA exactly restores the baseline (diff == 0) + +Observed on H100 (seed=42, 8 steps, 512x512): diff_1x~62.2, diff_2x~76.9, +diff_restored=0.0. Thresholds match test_bagel_expansion.py and hold with margin. +""" + +import json +import os +from pathlib import Path + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import numpy as np +import pytest +import torch +from PIL import Image +from safetensors.torch import save_file + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.lora.request import LoRARequest +from vllm_omni.lora.utils import stable_lora_int_id + +MODEL = "SenseNova/SenseNova-U1-8B-MoT" +DEFAULT_PROMPT = "a photo of a cat sitting on a laptop keyboard" +SIZE = (512, 512) # (width, height); small for fast E2E + +# SenseNova-U1-8B-MoT hidden_size (Qwen3-based). o_proj is square (hidden x hidden) +# regardless of GQA, so this is the robust LoRA target. Verify against config.json. +_LORA_DIM = 4096 +_LORA_MODULE = "language_model.model.layers.0.self_attn.o_proj_mot_gen" +_LORA_RANK = 4 + + +# --------------------------------------------------------------------------- +# Helpers (reused from test_sensenova_u1_text2img.py patterns) +# --------------------------------------------------------------------------- + + +def _build_sampling_params( + lora_request: LoRARequest | None = None, + lora_scale: float = 1.0, +) -> OmniDiffusionSamplingParams: + return OmniDiffusionSamplingParams( + height=SIZE[1], + width=SIZE[0], + seed=42, + num_inference_steps=8, + extra_args={ + "cfg_scale": 4.0, + "cfg_norm": "none", + "timestep_shift": 3.0, + "cfg_interval": (0.0, 1.0), + "batch_size": 1, + "think": False, + "t_eps": 0.02, + }, + lora_request=lora_request, + lora_scale=lora_scale, + ) + + +def _extract_generated_image(omni_outputs: list) -> Image.Image | None: + for req_output in omni_outputs: + if images := getattr(req_output, "images", None): + return images[0] + return None + + +def _generate(omni: Omni, lora_request: LoRARequest | None = None, lora_scale: float = 1.0) -> Image.Image: + outputs = list( + omni.generate( + prompts={"prompt": DEFAULT_PROMPT, "modalities": ["image"]}, + sampling_params_list=_build_sampling_params(lora_request, lora_scale), + ) + ) + img = _extract_generated_image(outputs) + assert img is not None, "No image generated" + return img + + +def _make_file_lora_request(adapter_dir: Path) -> LoRARequest: + """Write synthetic adapter to disk and return a file-backed LoRARequest.""" + adapter_dir.mkdir(parents=True, exist_ok=True) + gen = torch.Generator().manual_seed(42) + lora_a = torch.randn((_LORA_RANK, _LORA_DIM), dtype=torch.float32, generator=gen) * 0.1 + lora_b = torch.randn((_LORA_DIM, _LORA_RANK), dtype=torch.float32, generator=gen) * 0.5 + save_file( + { + f"base_model.model.{_LORA_MODULE}.lora_A.weight": lora_a, + f"base_model.model.{_LORA_MODULE}.lora_B.weight": lora_b, + }, + str(adapter_dir / "adapter_model.safetensors"), + ) + (adapter_dir / "adapter_config.json").write_text( + json.dumps({"r": _LORA_RANK, "lora_alpha": _LORA_RANK, "target_modules": [_LORA_MODULE]}), + encoding="utf-8", + ) + lora_dir = str(adapter_dir) + return LoRARequest(lora_name="test_file", lora_int_id=stable_lora_int_id(lora_dir), lora_path=lora_dir) + + +# --------------------------------------------------------------------------- +# Test +# --------------------------------------------------------------------------- + + +@pytest.mark.full_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None: + """Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation.""" + lora_request = _make_file_lora_request(tmp_path / "sensenova_lora") + + with OmniRunner(MODEL, stage_configs_path=None) as runner: + omni = runner.omni + baseline = _generate(omni) + img_1x = _generate(omni, lora_request, lora_scale=1.0) + img_2x = _generate(omni, lora_request, lora_scale=2.0) + restored = _generate(omni) + + baseline_arr = np.array(baseline, dtype=np.int16) + diff_1x = np.abs(baseline_arr - np.array(img_1x, dtype=np.int16)).mean() + diff_2x = np.abs(baseline_arr - np.array(img_2x, dtype=np.int16)).mean() + diff_restored = np.abs(baseline_arr - np.array(restored, dtype=np.int16)).mean() + + # (a) Adapter has visible effect at both scales + assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}" + assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}" + + # (b) Different scales produce different outputs + assert not np.isclose(diff_1x, diff_2x, atol=1.0), ( + f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}" + ) + + # (c) Output is not corrupted + assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}" + assert diff_2x < 120, f"LoRA output looks corrupted: diff_2x={diff_2x}" + + # (d) Deactivation fully restores base model + assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}" diff --git a/tests/e2e/offline_inference/test_voxcpm2.py b/tests/e2e/offline_inference/test_voxcpm2_tts.py similarity index 100% rename from tests/e2e/offline_inference/test_voxcpm2.py rename to tests/e2e/offline_inference/test_voxcpm2_tts.py diff --git a/tests/e2e/online_serving/test_higgs_audio_v3.py b/tests/e2e/online_serving/test_higgs_audio_v3.py new file mode 100644 index 000000000..aaa5498b3 --- /dev/null +++ b/tests/e2e/online_serving/test_higgs_audio_v3.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end online tests for higgs-audio v3 against /v1/audio/speech. + +Mirrors the higgs_audio_v2 test layout. Covers the plain-text-in / audio-out +happy path, a small concurrent burst (Stage 0 prefix caching has to coexist +with batching), the voice-clone path via ``ref_audio`` + ``ref_text``, the +upstream ``references[]`` cookbook alias added by +``normalize_references_alias`` in ``protocol/audio.py``, and the few +validator rejections that should remain 4xx. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") +os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0") + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import load_test_audio_data_url +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "bosonai/higgs-audio-v3-tts-4b" +STAGE_CONFIG = get_deploy_config_path("higgs_multimodal_qwen3.yaml") +SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"] +SERVER_ENV = {"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"} + +TEST_PARAMS = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIG, + server_args=SERVER_ARGS, + env_dict=SERVER_ENV, + ), + id="higgs_audio_v3_plain_text", + ) +] + +DEFAULT_SPEECH_TIMEOUT_S = 180.0 +# Floor for ~0.5 s of 24 kHz mono PCM_16: 24000 * 0.5 * 2 bytes ~= 24 KiB. +# A WAV header adds 44 bytes; conservative floor catches truncated / +# silence-only outputs without flagging short legitimate clips. +_MIN_AUDIO_BYTES = 20_000 + +# Reuse the shared TTS reference clip (clean ~5 s 24 kHz mono human speech) +# vendored under tests/assets/qwen3_tts/. Keeps a single WAV across TTS +# tests rather than duplicating asset bytes. +_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestHiggsAudioV3OnlineHappyPath: + """Plain-text -> audio happy paths against the live HTTP server.""" + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_plain_text_wav(self, omni_server, openai_client) -> None: + """Single non-streaming WAV request - canonical TTS happy path.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None: + """``max_new_tokens`` is one of the few extra fields the v3 validator accepts.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Innovation distinguishes between a leader and a follower.", + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "max_new_tokens": 500, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_concurrent_plain_text(self, omni_server, openai_client) -> None: + """Three concurrent non-streaming requests - guards per-slot audio state and Stage-0 PC under batching.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "It was the night before my birthday.", + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + }, + request_num=3, + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None: + """Streaming PCM via talker2code2wav_async_chunk + the connector's codec_streaming path. + + The deploy YAML pins ``async_chunk: true`` and ``codec_streaming: true`` so chunks + flow Stage 0 -> Stage 1 per AR loop. The Stage-1 codec honors + ``meta.left_context_size`` and ``meta.right_holdback_size`` so per-chunk codec + windows stitch into a coherent PCM stream. The byte-count gate is the same as + the sync paths; per-chunk audio content is verified offline against Whisper. + + NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the + check still catches catastrophic codec failure (silence, white noise, sample + scramble all give HNR << 0) while allowing for the sliding-window codec's + slightly-noisier-than-sync output. The default 1.0 dB threshold has only + ~0.16 dB margin over measured single-request output (1.16 dB on L4), so + flake risk is too high there. + """ + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The quick brown fox jumps over the lazy dog.", + "stream": True, + "response_format": "pcm", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + "min_hnr_db": 0.0, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None: + """Three concurrent streaming requests - guards per-request frame cursors + in ``talker2code2wav_async_chunk`` and per-slot delay-pattern state under batched AR. + + NOTE: ``min_hnr_db=-2.0`` is looser than the single-request streaming test + because batched AR adds per-request codec quality variance (DAC boundary + artifacts compounded across the 3-way batch); -2.0 dB gives ~1.6 dB margin + over the measured worst-of-3 on L4 (-0.42 dB) while still well above the + catastrophic-failure region. + """ + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "She sells seashells by the seashore.", + "stream": True, + "response_format": "pcm", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + "min_hnr_db": -2.0, + }, + request_num=3, + ) + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestHiggsAudioV3OnlineInlineControlTokens: + """Inline control-token surface from the bosonai cookbook. + + The model exposes four categories of ``<|category:value|>`` tokens + (emotion, style, prosody, sfx). Delivery tokens (emotion / style / + prosody speed-pitch-expressive) go at the START of input; positional + tokens (pause, sfx) go inline. SFX tokens must pair with their written + onomatopoeia. + + These tests check the *serving surface* - the validator accepts the + payload, the engine produces audio, and the WAV has non-trivial size. + They do not assert audio quality (out of scope for CI). + """ + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None: + """Delivery tokens (emotion + expressive_high) at the start of input.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": ("<|emotion:amusement|><|prosody:expressive_high|>Wait, that was actually hilarious."), + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_inline_style_whispering(self, omni_server, openai_client) -> None: + """Style token at the start - ``whispering``.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "<|style:whispering|>It is just between you and me, alright?", + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> None: + """Two prosody tokens at the start - slow speed plus low pitch.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": ( + "<|prosody:speed_slow|><|prosody:pitch_low|>The radar shows a storm approaching from the east." + ), + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_inline_pause_mid_text(self, omni_server, openai_client) -> None: + """Positional pause token placed inline between two clauses.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hold on a moment <|prosody:pause|> let me think about it.", + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_inline_sfx_with_onomatopoeia(self, omni_server, openai_client) -> None: + """SFX token paired with its written onomatopoeia (``<|sfx:laughter|>Hehe``).""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": ( + "<|emotion:amusement|>I cannot believe that just happened. " + "<|sfx:laughter|>Hehe, I am still recovering from it." + ), + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestHiggsAudioV3OnlineVoiceClone: + """Voice clone via the two payload shapes the model serves.""" + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_voice_clone_ref_audio_ref_text(self, omni_server, openai_client) -> None: + """Canonical vllm-omni voice clone via ``ref_audio`` + ``ref_text``.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "ref_audio": _REF_AUDIO_URL, + "ref_text": _REF_TEXT, + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_voice_clone_references_alias(self, omni_server, openai_client) -> None: + """BosonAI cookbook payload: ``references=[{audio_path, text}]``. + + ``normalize_references_alias`` (``protocol/audio.py``) is supposed to + translate the cookbook field into ``ref_audio`` / ``ref_text``; without + it the request would silently fall through to zero-shot synthesis. + """ + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "references": [{"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT}], + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestHiggsAudioV3OnlineValidatorRejections: + """Out-of-scope shapes must come back as 4xx. + + These cases use ``send_audio_speech_http_request`` (raw HTTP POST against + ``/v1/audio/speech``) rather than the OpenAI SDK helper because: + + - We need the failure-path matcher; the SDK helper asserts success. + - The cookbook ``references`` field is not in the OpenAI SDK schema, so + the SDK strips it before the request leaves the client and the server + sees a plain TTS request that returns 200. Raw HTTP keeps the JSON body + intact so the validator actually fires. + """ + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_empty_input(self, omni_server, openai_client) -> None: + """Empty text input - validator should reject before reaching the engine.""" + openai_client.send_audio_speech_http_request( + { + "json": { + "model": omni_server.model, + "input": "", + "response_format": "wav", + }, + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "err_code": (400, 422), + "err_message": ("input", "empty"), + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_multi_reference_payload(self, omni_server, openai_client) -> None: + """``references[]`` with more than one entry - multi-shot voice clone is not supported.""" + openai_client.send_audio_speech_http_request( + { + "json": { + "model": omni_server.model, + "input": "Hello world.", + "references": [ + {"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT}, + {"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT}, + ], + "response_format": "wav", + }, + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "err_code": (400, 422), + "err_message": "references", + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_conflicting_ref_audio_and_references(self, omni_server, openai_client) -> None: + """``ref_audio`` and ``references`` cannot both be set to different values.""" + openai_client.send_audio_speech_http_request( + { + "json": { + "model": omni_server.model, + "input": "Hello world.", + "ref_audio": _REF_AUDIO_URL, + "references": [{"audio_path": "https://example.com/other.wav"}], + "response_format": "wav", + }, + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "err_code": (400, 422), + "err_message": "mutually exclusive", + } + ) diff --git a/tests/e2e/online_serving/test_ming_tts.py b/tests/e2e/online_serving/test_ming_tts.py new file mode 100644 index 000000000..03982cc4a --- /dev/null +++ b/tests/e2e/online_serving/test_ming_tts.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E online serving tests for Ming-omni-tts (dense) model. +Tests text-to-audio via /v1/audio/speech endpoint. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +pytestmark = [pytest.mark.advanced_model, pytest.mark.omni] + +MODEL = "inclusionAI/Ming-omni-tts-0.5B" +DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml") + +no_async_chunk_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=DEPLOY_CONFIG, + server_args=["--enforce-eager", "--no-async-chunk"], + ), + id="no_async_chunk", + ) +] + +async_chunk_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=DEPLOY_CONFIG, + server_args=["--enforce-eager"], + ), + id="async_chunk", + ) +] + + +def get_prompt(prompt_type="zh"): + prompts = { + "zh": "今天天气真不错，适合出去散散步。", + "zh_short": "这款产品的名字，叫变态坑爹牛肉丸。", + } + return prompts.get(prompt_type, prompts["zh"]) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", no_async_chunk_params, indirect=True) +def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None: + """ + Deploy Setting: ming_tts.yaml with --no-async-chunk + Input Modal: text + Output Modal: audio + Input Setting: stream=False + Datasets: two concurrent requests + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("zh"), + "stream": False, + "response_format": "wav", + "timeout": 300.0, + } + openai_client.send_audio_speech_request(request_config, request_num=2) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", async_chunk_params, indirect=True) +def test_text_to_audio_streaming_001(omni_server, openai_client) -> None: + """ + Deploy Setting: ming_tts.yaml (async_chunk=true) + Input Modal: text + voice + Output Modal: audio (streamed) + Input Setting: stream=True + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("zh_short"), + "voice": "灵小甄", + "stream": True, + "response_format": "wav", + "timeout": 300.0, + } + openai_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_voxtral_tts.py b/tests/e2e/online_serving/test_voxtral_tts.py index a1b5fdbb8..a86de2100 100644 --- a/tests/e2e/online_serving/test_voxtral_tts.py +++ b/tests/e2e/online_serving/test_voxtral_tts.py @@ -91,3 +91,90 @@ class TestVoxtralTTSFixedVoice: "timeout": 120.0, } ) + + @pytest.mark.advanced_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_speech_speed(self, omni_server, openai_client) -> None: + """Request with speed parameters""" + speeds = [0.5, 1, 1.5, 2, 2.5] + for speed in speeds: + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The boy was there when the sun rose.", + "voice": "casual_female", + "language": "English", + "response_format": "wav", + "timeout": 120.0, + "speed": speed, + } + ) + + @pytest.mark.advanced_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_speech_instructions(self, omni_server, openai_client) -> None: + """Request with instructions parameters""" + instructions = [ + "Speak formally", + "Speak angrily", + "Deliver with a sad voice", + "Speak with a chirpy happy voice", + ] + for instruction in instructions: + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The boy was there when the sun rose.", + "voice": "casual_female", + "language": "English", + "response_format": "wav", + "timeout": 120.0, + "instructions": instruction, + } + ) + + @pytest.mark.advanced_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_speech_response_formats(self, omni_server, openai_client) -> None: + """Test TTS with different response formats""" + response_formats = ["wav", "mp3"] + for response_format in response_formats: + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Testing various response formats.", + "voice": "casual_male", + "language": "English", + "response_format": response_format, + "timeout": 120.0, + } + ) + + @pytest.mark.advanced_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_speech_batches(self, omni_server, openai_client) -> None: + """Test TTS batches""" + items = [ + {"input": "The birch canoe slid on the smooth planks."}, + {"input": "Glue the sheet to the dark blue background."}, + {"input": "It's easy to tell the depth of a well."}, + {"input": "These days a chicken leg is a rare dish."}, + {"input": "Rice is often served in round bowls."}, + ] + + openai_client.send_audio_speech_batch_http_request( + { + "json": { + "model": omni_server.model, + "items": items, + "voice": "casual_male", + "language": "English", + "response_format": "wav", + }, + "timeout": 120.0, + } + ) diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index c5cba1ea9..0553c03ff 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -161,3 +161,113 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) - finally: orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_diffusion_client_error_output_propagates_status_code(orchestrator_factory) -> None: + """A client-error OmniRequestOutput (e.g. a guardrail 4xx) from a diffusion + stage must surface as an ErrorMessage that carries status_code/error_type, + so the frontend can map it to the correct HTTP response instead of a 500. + """ + stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image") + orchestrator_fixture = orchestrator_factory([stage0]) + params = OmniDiffusionSamplingParams() + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-blocked", + prompt={"prompt": "blocked"}, + original_prompt={"prompt": "blocked"}, + sampling_params_list=[params], + final_stage_id=0, + ) + + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + + # Push a client-error output from the diffusion stage. + stage0.push_diffusion_output( + OmniRequestOutput.from_error( + "req-blocked", + "Input was blocked by Cosmos3 guardrails.", + status_code=400, + error_type="BadRequestError", + ) + ) + + msg = await _get_any_output_message(orchestrator_fixture) + + assert isinstance(msg, ErrorMessage) + assert msg.request_id == "req-blocked" + assert msg.stage_id == 0 + assert msg.fatal is False + assert msg.status_code == 400 + assert msg.error_type == "BadRequestError" + assert msg.error == "Input was blocked by Cosmos3 guardrails." + + # Request state should be cleaned up. + await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states) + finally: + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "error_type"), + [ + pytest.param(429, "RateLimitError", id="429-too-many-requests"), + pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"), + pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"), + pytest.param(403, "PermissionDeniedError", id="403-forbidden"), + ], +) +async def test_diffusion_client_error_output_propagates_non_400_status( + orchestrator_factory, + status_code: int, + error_type: str, +) -> None: + """A diffusion-stage client error with a non-400 4xx status must be serialized + into an ErrorMessage that preserves that exact status_code/error_type. + """ + stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image") + orchestrator_fixture = orchestrator_factory([stage0]) + params = OmniDiffusionSamplingParams() + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-blocked", + prompt={"prompt": "blocked"}, + original_prompt={"prompt": "blocked"}, + sampling_params_list=[params], + final_stage_id=0, + ) + + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + + # Push a non-400 client-error output from the diffusion stage. + stage0.push_diffusion_output( + OmniRequestOutput.from_error( + "req-blocked", + "client error from diffusion stage", + status_code=status_code, + error_type=error_type, + ) + ) + + msg = await _get_any_output_message(orchestrator_fixture) + + assert isinstance(msg, ErrorMessage) + assert msg.request_id == "req-blocked" + assert msg.stage_id == 0 + assert msg.fatal is False + assert msg.status_code == status_code + assert msg.error_type == error_type + assert msg.error == "client error from diffusion stage" + + # Request state should be cleaned up. + await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states) + finally: + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) diff --git a/tests/entrypoints/openai_api/test_decode_audio_url.py b/tests/entrypoints/openai_api/test_decode_audio_url.py new file mode 100644 index 000000000..2a5c5be47 --- /dev/null +++ b/tests/entrypoints/openai_api/test_decode_audio_url.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for decode_audio_url helper.""" + +import asyncio +import base64 +import os + +import httpx +import pytest + +from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError +from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _run(coro): + return asyncio.run(coro) + + +def _cleanup(path): + if path and os.path.exists(path): + os.unlink(path) + + +class TestDecodeAudioUrlDataUrl: + """Tests for base64 data URL decoding.""" + + def test_valid_mp3_data_url(self): + audio_bytes = b"\xff\xfb\x90\x00" * 100 + b64 = base64.b64encode(audio_bytes).decode() + url = f"data:audio/mp3;base64,{b64}" + + path = _run(decode_audio_url(url)) + try: + assert os.path.isfile(path) + assert path.endswith(".mp3") + with open(path, "rb") as f: + assert f.read() == audio_bytes + finally: + _cleanup(path) + + def test_valid_wav_data_url(self): + audio_bytes = b"RIFF" + b"\x00" * 100 + b64 = base64.b64encode(audio_bytes).decode() + url = f"data:audio/wav;base64,{b64}" + + path = _run(decode_audio_url(url)) + try: + assert path.endswith(".wav") + with open(path, "rb") as f: + assert f.read() == audio_bytes + finally: + _cleanup(path) + + def test_valid_mpeg_data_url_uses_mp3_suffix(self): + audio_bytes = b"\xff\xfb\x90\x00" * 50 + b64 = base64.b64encode(audio_bytes).decode() + url = f"data:audio/mpeg;base64,{b64}" + + path = _run(decode_audio_url(url)) + try: + assert path.endswith(".mp3") + finally: + _cleanup(path) + + def test_valid_ogg_data_url(self): + audio_bytes = b"OggS" + b"\x00" * 100 + b64 = base64.b64encode(audio_bytes).decode() + url = f"data:audio/ogg;base64,{b64}" + + path = _run(decode_audio_url(url)) + try: + assert path.endswith(".ogg") + finally: + _cleanup(path) + + def test_invalid_base64_raises(self): + url = "data:audio/mp3;base64,!!!not-valid-base64!!!" + + with pytest.raises(InvalidInputReferenceError, match="not valid base64"): + _run(decode_audio_url(url)) + + def test_empty_audio_data_raises(self): + url = "data:audio/mp3;base64," + + with pytest.raises(InvalidInputReferenceError, match="audio data is empty"): + _run(decode_audio_url(url)) + + +class TestDecodeAudioUrlSuffixSanitization: + """Tests for MIME extension sanitization (path traversal prevention).""" + + def test_path_traversal_in_mime_is_neutralized(self): + audio_bytes = b"\x00" * 100 + b64 = base64.b64encode(audio_bytes).decode() + url = f"data:audio/../..;base64,{b64}" + + path = _run(decode_audio_url(url)) + try: + # ".." contains non-alnum dots, so falls back to .wav + assert path.endswith(".wav") + assert os.path.dirname(path) == "/tmp" + finally: + _cleanup(path) + + def test_non_alnum_extension_falls_back_to_wav(self): + audio_bytes = b"\x00" * 100 + b64 = base64.b64encode(audio_bytes).decode() + url = f"data:audio/foo-bar;base64,{b64}" + + path = _run(decode_audio_url(url)) + try: + assert path.endswith(".wav") + finally: + _cleanup(path) + + def test_too_long_extension_falls_back_to_wav(self): + audio_bytes = b"\x00" * 100 + b64 = base64.b64encode(audio_bytes).decode() + url = f"data:audio/verylongextension;base64,{b64}" + + path = _run(decode_audio_url(url)) + try: + assert path.endswith(".wav") + finally: + _cleanup(path) + + def test_valid_short_alnum_extension_is_used(self): + audio_bytes = b"\x00" * 100 + b64 = base64.b64encode(audio_bytes).decode() + url = f"data:audio/flac;base64,{b64}" + + path = _run(decode_audio_url(url)) + try: + assert path.endswith(".flac") + finally: + _cleanup(path) + + +class TestDecodeAudioUrlHttp: + """Tests for HTTP URL decoding.""" + + def test_valid_http_url(self, monkeypatch): + audio_bytes = b"\xff\xfb\x90\x00" * 50 + fake_resp = httpx.Response(200, content=audio_bytes) + + async def _mock_get(self, url): + return fake_resp + + monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get) + monkeypatch.setattr(httpx.Response, "raise_for_status", lambda self: None) + + path = _run(decode_audio_url("https://example.com/audio.mp3")) + try: + assert os.path.isfile(path) + assert path.endswith(".wav") + with open(path, "rb") as f: + assert f.read() == audio_bytes + finally: + _cleanup(path) + + def test_http_error_raises(self, monkeypatch): + async def _mock_get(self, url): + raise httpx.HTTPStatusError("Not Found", request=httpx.Request("GET", url), response=httpx.Response(404)) + + monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get) + + with pytest.raises(InvalidInputReferenceError, match="failed to download audio"): + _run(decode_audio_url("https://example.com/missing.mp3")) + + +class TestDecodeAudioUrlInvalid: + """Tests for invalid URL schemes.""" + + def test_local_file_path_raises(self): + with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"): + _run(decode_audio_url("/path/to/local/file.mp3")) + + def test_ftp_url_raises(self): + with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"): + _run(decode_audio_url("ftp://example.com/file.mp3")) + + def test_empty_string_raises(self): + with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"): + _run(decode_audio_url("")) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 38cb86c6e..d108b9323 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -29,6 +29,7 @@ from vllm_omni.entrypoints.openai.image_api_utils import ( parse_size, ) from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.errors import GuardrailViolationError from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -587,6 +588,24 @@ def test_generate_single_image(test_client): assert test_client.app.state.engine_client.captured_prompt["modalities"] == ["image"] +def test_generate_images_guardrail_error_returns_400(test_client, mock_async_diffusion): + async def blocked_generate(**kwargs): + raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.") + yield MockGenerationResult([]) # pragma: no cover + + mock_async_diffusion.generate = blocked_generate + response = test_client.post( + "/v1/images/generations", + json={ + "prompt": "blocked prompt", + "n": 1, + "size": "1024x1024", + }, + ) + assert response.status_code == 400 + assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails." + + def test_generate_images_async_omni_sampling_params(async_omni_test_client): """Test AsyncOmni path uses per-stage sampling params.""" response = async_omni_test_client.post( @@ -891,6 +910,31 @@ def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_ assert payloads[2] == "[DONE]" +def test_image_edits_streaming_guardrail_error_uses_400(streaming_image_edit_client): + async def generate_guardrail_error(prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs): + raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.") + yield MockStageResult(stage_id=1, final_output_type="image", images=[]) # pragma: no cover + + streaming_image_edit_client.app.state.engine_client.generate = generate_guardrail_error + img_bytes = make_test_image_bytes((16, 16)) + response = streaming_image_edit_client.post( + "/v1/images/edits", + files=[("image", img_bytes)], + data={ + "prompt": "blocked prompt", + "stream": "true", + }, + ) + + assert response.status_code == 200 + payloads = _parse_sse_payloads(response.text) + assert payloads[0]["object"] == "error" + assert payloads[0]["error"]["message"] == "Input was blocked by Cosmos3 guardrails." + assert payloads[0]["error"]["type"] == "BadRequestError" + assert payloads[0]["error"]["code"] == 400 + assert payloads[1] == "[DONE]" + + def test_image_edits_streaming_rejects_single_stage(test_client): img_bytes = make_test_image_bytes((16, 16)) response = test_client.post( diff --git a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py index c8bb050e7..507018ef7 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py +++ b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py @@ -61,7 +61,7 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser assert engine_prompt["prompt"] == "draw a robot" assert engine_prompt["modalities"] == ["img2img"] assert engine_prompt["negative_prompt"] == "blurry" - assert engine_prompt["mm_processor_kwargs"] == {"target_h": 768, "target_w": 1024} + assert engine_prompt["mm_processor_kwargs"] == {"target_h": 768, "target_w": 1024, "vae_generator_seed": 0} assert engine_prompt["multi_modal_data"]["img2img"].size == (24, 24) assert len(sampling_params_list) == 3 diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 4fcf3418a..d4feeef60 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -31,8 +31,10 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( CreateAudio, OpenAICreateSpeechRequest, SpeechBatchItem, + StreamingSpeechSessionConfig, ) from vllm_omni.entrypoints.openai.serving_speech import ( + _TTS_LANGUAGES, OmniOpenAIServingSpeech, _create_wav_header, ) @@ -1234,6 +1236,49 @@ class TestTTSMethods: assert params["language"] == ["English"] assert params["task_type"] == ["CustomVoice"] + def test_build_tts_params_base_non_streaming_mode_true(self, speech_server): + """Base task should pass through an explicit non_streaming_mode override.""" + req = OpenAICreateSpeechRequest( + input="Hello", + task_type="Base", + ref_audio="data:audio/wav;base64,abc", + ref_text="reference", + non_streaming_mode=True, + ) + + params = speech_server._build_tts_params(req) + + assert params["task_type"] == ["Base"] + assert params["non_streaming_mode"] == [True] + + def test_build_tts_params_base_omits_non_streaming_mode_by_default(self, speech_server): + """Base task should keep using the model default when no override is sent.""" + req = OpenAICreateSpeechRequest( + input="Hello", + task_type="Base", + ref_audio="data:audio/wav;base64,abc", + ref_text="reference", + ) + + params = speech_server._build_tts_params(req) + + assert params["task_type"] == ["Base"] + assert "non_streaming_mode" not in params + + def test_build_tts_params_explicit_non_streaming_mode_overrides_voicedesign_default(self, speech_server): + """Explicit false should not be replaced by the VoiceDesign fallback.""" + req = OpenAICreateSpeechRequest( + input="Hello", + task_type="VoiceDesign", + instructions="warm and calm", + non_streaming_mode=False, + ) + + params = speech_server._build_tts_params(req) + + assert params["task_type"] == ["VoiceDesign"] + assert params["non_streaming_mode"] == [False] + def test_load_supported_speakers(self, mocker: MockerFixture): """Test _load_supported_speakers.""" mock_engine_client = mocker.MagicMock() @@ -1257,6 +1302,67 @@ class TestTTSMethods: # Verify speakers are normalized to lowercase assert server.supported_speakers == {"ryan", "vivian", "aiden"} + def test_load_supported_languages_from_config(self, speech_server): + """Languages/dialects from codec_language_id are loaded title-cased; 'Auto' is added.""" + speech_server._tts_model_type = "qwen3_tts" + speech_server.engine_client.model_config = SimpleNamespace( + hf_config=SimpleNamespace( + talker_config=SimpleNamespace( + codec_language_id={"chinese": 2055, "english": 2050, "beijing_dialect": 2074} + ) + ) + ) + assert speech_server._load_supported_languages() == { + "Chinese", + "English", + "Beijing_Dialect", + "Auto", + } + + def test_load_supported_languages_from_dict_config(self, speech_server): + """talker_config provided as a plain dict is handled the same as an object.""" + speech_server._tts_model_type = "qwen3_tts" + speech_server.engine_client.model_config = SimpleNamespace( + hf_config=SimpleNamespace(talker_config={"codec_language_id": {"chinese": 2055, "english": 2050}}) + ) + assert speech_server._load_supported_languages() == {"Chinese", "English", "Auto"} + + def test_validate_language_custom_dialect_accepted(self, speech_server): + """A language present in the model config passes validation, case-insensitively.""" + speech_server.supported_languages = {"Chinese", "English", "Beijing_Dialect", "Auto"} + for language in ("Beijing_Dialect", "beijing_dialect", "English", "english", "Auto", "AUTO"): + req = OpenAICreateSpeechRequest(input="Hello", language=language) + result = speech_server._validate_tts_request(req) + assert result is None or "Invalid language" not in result + # Language is normalized to the title-cased config form. + assert req.language == language.title() + + def test_validate_language_unknown_rejected(self, speech_server): + """A language not in the configured set is rejected.""" + speech_server.supported_languages = {"Chinese", "English", "Auto"} + for language in ("Klingon", "klingon"): + req = OpenAICreateSpeechRequest(input="Hello", language=language) + assert "Invalid language" in speech_server._validate_tts_request(req) + + def test_load_supported_languages_default_when_no_config(self, speech_server): + """Empty/missing codec_language_id on a Qwen3-TTS model falls back to the default list.""" + speech_server._tts_model_type = "qwen3_tts" + speech_server.engine_client.model_config = SimpleNamespace( + hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_language_id={})) + ) + assert speech_server._load_supported_languages() == _TTS_LANGUAGES + + def test_load_supported_languages_default_on_config_error(self, speech_server): + """If the model config cannot be read, fall back to the default list.""" + speech_server._tts_model_type = "qwen3_tts" + speech_server.engine_client = SimpleNamespace() # no model_config -> AttributeError + assert speech_server._load_supported_languages() == _TTS_LANGUAGES + + def test_load_supported_languages_default_for_non_qwen(self, speech_server): + """Non-qwen3_tts model types get the default language set.""" + speech_server._tts_model_type = None + assert speech_server._load_supported_languages() == _TTS_LANGUAGES + def test_build_tts_params_with_uploaded_voice(self, speech_server, mocker: MockerFixture): """Test _build_tts_params auto-sets ref_audio for uploaded voices (x_vector only).""" speech_server.uploaded_speakers = { @@ -1961,6 +2067,34 @@ class TestMergeBatchItem: assert merged.task_type == "CustomVoice" assert merged.max_new_tokens == 512 + def test_non_streaming_mode_batch_default_used(self): + """Batch-level non_streaming_mode should be used when item doesn't specify one.""" + batch = BatchSpeechRequest( + items=[SpeechBatchItem(input="hi")], + non_streaming_mode=True, + ) + + merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0]) + + assert merged.non_streaming_mode is True + + def test_non_streaming_mode_item_override_wins(self): + """Per-item false should override a true batch-level default.""" + batch = BatchSpeechRequest( + items=[SpeechBatchItem(input="hi", non_streaming_mode=False)], + non_streaming_mode=True, + ) + + merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0]) + + assert merged.non_streaming_mode is False + + +def test_streaming_speech_session_config_accepts_non_streaming_mode(): + config = StreamingSpeechSessionConfig(non_streaming_mode=True) + + assert config.non_streaming_mode is True + class TestAsyncOmniSupportedTasks: """Test that AsyncOmni reports correct supported tasks based on output modalities.""" @@ -2749,6 +2883,52 @@ class TestTTSAsyncOffloading: qwen3_tts_server._build_tts_params.assert_called_once() qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once() + def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false( + self, qwen3_tts_server, mocker: MockerFixture + ): + """VoiceDesign explicit false should reach the model prompt additional_information.""" + qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + + request = OpenAICreateSpeechRequest( + input="hello", + task_type="VoiceDesign", + instructions="warm and calm", + non_streaming_mode=False, + ) + _request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + assert tts_params["task_type"] == ["VoiceDesign"] + assert tts_params["non_streaming_mode"] == [False] + prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"] + assert prompt["additional_information"] is tts_params + assert prompt["additional_information"]["non_streaming_mode"] == [False] + + def test_prepare_speech_generation_qwen3_base_non_streaming_mode_true( + self, qwen3_tts_server, mocker: MockerFixture + ): + """Base explicit true should reach the model prompt additional_information.""" + qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) + qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000)) + qwen3_tts_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value=None) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + + request = OpenAICreateSpeechRequest( + input="hello", + task_type="Base", + ref_audio="data:audio/wav;base64,abc", + ref_text="reference transcript", + non_streaming_mode=True, + ) + _request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + assert tts_params["task_type"] == ["Base"] + assert tts_params["ref_text"] == ["reference transcript"] + assert tts_params["non_streaming_mode"] == [True] + prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"] + assert prompt["additional_information"] is tts_params + assert prompt["additional_information"]["non_streaming_mode"] == [True] + def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self, qwen3_tts_server): """After a ref artifact is marked ready, repeated requests avoid ref_audio payload IPC.""" wav_list = [0.0] * 48000 diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 072fd25f1..b54afd54a 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -13,12 +13,14 @@ import threading import time from types import SimpleNamespace +import numpy as np import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from PIL import Image from pytest_mock import MockerFixture +from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes from vllm_omni.entrypoints.openai import api_server from vllm_omni.entrypoints.openai.api_server import router from vllm_omni.entrypoints.openai.protocol.videos import ( @@ -29,6 +31,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo from vllm_omni.entrypoints.openai.storage import LocalStorageManager from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry +from vllm_omni.errors import GuardrailViolationError from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -84,7 +87,10 @@ class BlockingVideoHandler: if self.stage_configs is None: self.stage_configs = stage_configs - async def generate_video_bytes(self, request, reference_id, *, reference_image=None): + async def generate_video_bytes( + self, request, reference_id, *, reference_image=None, reference_video=None, reference_audio=None + ): + del request, reference_id, reference_image, reference_video, reference_audio self.started.set() try: await asyncio.Future() @@ -130,6 +136,30 @@ def _make_test_image_data_url(size=(64, 64)) -> str: return f"data:image/png;base64,{encoded}" +def _make_test_video_bytes(size=(32, 24), num_frames=3) -> bytes: + width, height = size + frames = np.zeros((num_frames, height, width, 3), dtype=np.uint8) + for idx in range(num_frames): + frames[idx, :, :, 0] = idx * 40 + frames[idx, :, :, 1] = 128 + frames[idx, :, :, 2] = 255 - idx * 40 + return mux_video_audio_bytes(frames, fps=8, video_codec_options={"preset": "ultrafast", "threads": "0"}) + + +def _make_test_video_data_url(size=(32, 24), num_frames=3) -> str: + encoded = base64.b64encode(_make_test_video_bytes(size, num_frames)).decode("utf-8") + return f"data:video/mp4;base64,{encoded}" + + +def _cosmos3_stage_configs(): + return [ + SimpleNamespace( + stage_type="diffusion", + engine_args=SimpleNamespace(model_class_name="Cosmos3OmniDiffusersPipeline"), + ) + ] + + def _wait_for_status(client: TestClient, video_id: str, status: str, timeout_s: float = 2.0): deadline = time.time() + timeout_s last_payload = None @@ -332,6 +362,139 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc assert input_image.size == (40, 24) +def test_v2v_video_generation_form(test_client, mocker: MockerFixture): + video_bytes = _make_test_video_bytes((32, 24), num_frames=3) + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + response = test_client.post( + "/v1/videos", + data={"prompt": "Continue this motion."}, + files={"input_reference": ("input.mp4", video_bytes, "video/mp4")}, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + engine = test_client.app.state.openai_serving_video._engine_client + prompt = engine.captured_prompt + assert "multi_modal_data" in prompt + assert "video" in prompt["multi_modal_data"] + input_video = prompt["multi_modal_data"]["video"] + assert len(input_video) == 3 + assert all(isinstance(frame, Image.Image) for frame in input_video) + assert input_video[0].size == (32, 24) + + +def test_v2v_video_generation_with_video_reference_form(test_client, mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "Continue this motion.", + "video_reference": json.dumps({"video_url": _make_test_video_data_url((32, 24), 2)}), + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + engine = test_client.app.state.openai_serving_video._engine_client + input_video = engine.captured_prompt["multi_modal_data"]["video"] + assert len(input_video) == 2 + assert input_video[0].size == (32, 24) + + +def test_decode_video_bytes_can_keep_last_frames(): + from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes + + frames = _decode_video_bytes( + _make_test_video_bytes((32, 24), num_frames=6), + source="input_reference", + max_frames=2, + keep="last", + ) + + assert len(frames) == 2 + red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames] + assert red_means[0] > 100 + assert red_means[1] > red_means[0] + + +def test_cosmos3_reference_video_limit_uses_v2v_condition_frames(): + request = VideoGenerationRequest( + prompt="Continue this motion.", + num_frames=189, + extra_params={"condition_frame_indexes_vision": [0, 2]}, + ) + + spec = api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()) + assert spec.max_frames == 9 + assert spec.keep == "first" + + +def test_cosmos3_reference_video_limit_preserves_action_frames(): + request = VideoGenerationRequest( + prompt="Predict the action.", + num_frames=17, + extra_params={"action_mode": "inverse_dynamics", "action_chunk_size": 16}, + ) + + assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 17 + + +def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames(): + request = VideoGenerationRequest( + prompt="Continue this motion.", + num_frames=5, + extra_params={"condition_frame_indexes_vision": [0, 20]}, + ) + + assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5 + + +def test_s2v_video_generation_with_audio_reference_form(test_client, mocker: MockerFixture): + """Speech-to-video: image + audio_reference (base64 data URL) passes audio path to multi_modal_data.""" + audio_bytes = b"\xff\xfb\x90\x00" * 50 + audio_b64 = base64.b64encode(audio_bytes).decode() + audio_ref = json.dumps({"audio_url": f"data:audio/mp3;base64,{audio_b64}"}) + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "A person singing", + "audio_reference": audio_ref, + "width": "832", + "height": "480", + }, + files={"input_reference": ("face.png", _make_test_image_bytes((64, 64)), "image/png")}, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + engine = test_client.app.state.openai_serving_video._engine_client + prompt = engine.captured_prompt + assert "multi_modal_data" in prompt + assert "image" in prompt["multi_modal_data"] + assert "audio" in prompt["multi_modal_data"] + audio_path = prompt["multi_modal_data"]["audio"] + assert isinstance(audio_path, str) + assert audio_path.endswith(".mp3") + + def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture): fps_values = [] @@ -794,7 +957,20 @@ def test_rejects_input_reference_and_image_reference_together(test_client): files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")}, ) assert response.status_code == 400 - assert "either input_reference or image_reference" in response.json()["detail"].lower() + assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower() + + +def test_rejects_image_reference_and_video_reference_together(test_client): + response = test_client.post( + "/v1/videos", + data={ + "prompt": "bad refs", + "image_reference": '{"image_url": "https://example.com/cat.png"}', + "video_reference": '{"video_url": "https://example.com/cat.mp4"}', + }, + ) + assert response.status_code == 400 + assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower() def test_invalid_seconds_returns_422(test_client): @@ -844,6 +1020,25 @@ def test_invalid_lora_returns_400(test_client): assert "lora object" in failed["error"]["message"].lower() +def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture): + mocker.patch.object( + OmniOpenAIServingVideo, + "generate_video_bytes", + side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."), + ) + response = test_client.post("/v1/videos", data={"prompt": "blocked prompt"}) + assert response.status_code == 200 + + video_id = response.json()["id"] + failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value) + assert failed["error"]["code"] == 400 + assert failed["error"]["message"] == "Input was blocked by Cosmos3 guardrails." + + retrieve = test_client.get(f"/v1/videos/{video_id}") + assert retrieve.status_code == 400 + assert retrieve.json()["error"]["code"] == 400 + + def test_unsupported_image_reference_file_id_returns_400(test_client): response = test_client.post( "/v1/videos", @@ -856,6 +1051,18 @@ def test_unsupported_image_reference_file_id_returns_400(test_client): assert response.json()["detail"] == "Invalid image_reference: file_id is not supported yet." +def test_unsupported_video_reference_file_id_returns_400(test_client): + response = test_client.post( + "/v1/videos", + data={ + "prompt": "unsupported ref", + "video_reference": '{"file_id": "file-123"}', + }, + ) + assert response.status_code == 400 + assert response.json()["detail"] == "Invalid video_reference: file_id is not supported yet." + + def test_invalid_uploaded_input_reference_returns_400(test_client): response = test_client.post( "/v1/videos", @@ -863,7 +1070,7 @@ def test_invalid_uploaded_input_reference_returns_400(test_client): files={"input_reference": ("input.png", b"not-an-image", "image/png")}, ) assert response.status_code == 400 - assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image." + assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image or video." def test_video_request_validation(): @@ -880,6 +1087,8 @@ def test_video_request_validation(): with pytest.raises(ValueError): VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"}) + with pytest.raises(ValueError): + VideoGenerationRequest(prompt="test", video_reference={"file_id": "file-1", "video_url": "https://example.com"}) with pytest.raises(ValueError): VideoGenerationRequest(prompt="test", frame_interpolation_exp=0) with pytest.raises(ValueError): @@ -1265,6 +1474,23 @@ def test_sync_i2v_with_image_reference(test_client, mocker: MockerFixture): assert response.content == b"ref-video" +def test_sync_v2v_returns_video_bytes(test_client, mocker: MockerFixture): + video_bytes = _make_test_video_bytes((32, 24), num_frames=3) + _mock_encode_video_bytes(mocker, b"v2v-video-data") + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "Continue this motion."}, + files={"input_reference": ("input.mp4", video_bytes, "video/mp4")}, + ) + + assert response.status_code == 200 + assert response.content == b"v2v-video-data" + engine = test_client.app.state.openai_serving_video._engine_client + input_video = engine.captured_prompt["multi_modal_data"]["video"] + assert len(input_video) == 3 + assert input_video[0].size == (32, 24) + + def test_sync_missing_handler_returns_503(): app = FastAPI() app.include_router(router) @@ -1297,7 +1523,7 @@ def test_sync_rejects_both_references(test_client): files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")}, ) assert response.status_code == 400 - assert "either input_reference or image_reference" in response.json()["detail"].lower() + assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower() def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture): @@ -1315,6 +1541,20 @@ def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture): assert "GPU exploded" in response.json()["detail"] +def test_sync_guardrail_error_returns_400(test_client, mocker: MockerFixture): + mocker.patch.object( + OmniOpenAIServingVideo, + "generate_video_bytes", + side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."), + ) + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "blocked prompt"}, + ) + assert response.status_code == 400 + assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails." + + def test_sync_does_not_create_store_entry(test_client, mocker: MockerFixture): """The sync endpoint should NOT leave any record in VIDEO_STORE — it is stateless by design.""" diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 28717d28c..8f647138b 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -18,6 +18,12 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni import Omni from vllm_omni.entrypoints.omni_base import OmniEngineDeadError +from vllm_omni.errors import ( + OmniClientError, + client_error_from_metadata, + client_error_metadata, + is_client_error_status, +) from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -604,6 +610,71 @@ async def test_async_omni_propagates_fatal_error_context(monkeypatch: pytest.Mon assert getattr(exc_info.value, "error_stage_id") == 2 +def _enqueue_client_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None: + engine.output_q.put_nowait( + ErrorMessage( + request_id=msg["request_id"], + stage_id=2, + error="Input was blocked by Cosmos3 guardrails.", + status_code=400, + error_type="BadRequestError", + ) + ) + + +@pytest.mark.asyncio +async def test_async_omni_propagates_client_error_status(monkeypatch: pytest.MonkeyPatch): + """A non-fatal client error from the orchestrator must be routed to the + requesting generate() call (not raised in the shared dispatcher) and + surface as an OmniClientError carrying status_code/error_type.""" + engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_client_error_message) + _patch_engine(monkeypatch, engine) + + app = AsyncOmni("dummy-model") + try: + with pytest.raises(OmniClientError) as exc_info: + async for _ in app.generate(prompt="blocked", request_id="req-1"): + pass + finally: + app.shutdown() + + assert exc_info.value.status_code == 400 + assert exc_info.value.error_type == "BadRequestError" + assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails." + + +def _enqueue_server_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None: + engine.output_q.put_nowait( + ErrorMessage( + request_id=msg["request_id"], + stage_id=2, + error="GPU exploded", + ) + ) + + +@pytest.mark.asyncio +async def test_async_omni_propagates_server_error_as_runtime(monkeypatch: pytest.MonkeyPatch): + """A non-fatal error WITHOUT a 4xx status_code is a server fault: it must + surface as a RuntimeError (-> HTTP 500), not an OmniClientError (-> 400). + Guards against the fix over-broadly mapping every error to a client error.""" + engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_server_error_message) + _patch_engine(monkeypatch, engine) + + app = AsyncOmni("dummy-model") + try: + # OmniClientError subclasses ValueError, not RuntimeError, so matching + # RuntimeError here also proves it was NOT raised as a client error. + with pytest.raises(RuntimeError) as exc_info: + async for _ in app.generate(prompt="boom", request_id="req-1"): + pass + finally: + app.shutdown() + + assert not isinstance(exc_info.value, OmniClientError) + assert str(exc_info.value) == "GPU exploded" + + def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch): sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)] engine = FakeAsyncOmniEngine( @@ -761,6 +832,140 @@ def test_non_fatal_error_raises_runtime(): base._handle_output_message(msg) +def test_non_fatal_client_error_raises_omni_client_error(): + """A non-fatal ErrorMessage carrying a 4xx status_code (e.g. a guardrail + block) must surface as an OmniClientError with the metadata intact, not a + bare RuntimeError. This covers the offline/sync Omni consumer path.""" + base = _make_base() + msg = ErrorMessage( + error="Input was blocked by Cosmos3 guardrails.", + status_code=400, + error_type="BadRequestError", + ) + + with pytest.raises(OmniClientError) as exc_info: + base._handle_output_message(msg) + + assert exc_info.value.status_code == 400 + assert exc_info.value.error_type == "BadRequestError" + assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails." + + +_NON_400_CLIENT_ERRORS = [ + pytest.param(429, "RateLimitError", id="429-too-many-requests"), + pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"), + pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"), + pytest.param(403, "PermissionDeniedError", id="403-forbidden"), +] + + +@pytest.mark.parametrize( + ("status_code", "error_type"), [pytest.param(400, "BadRequestError", id="400")] + _NON_400_CLIENT_ERRORS +) +def test_client_error_metadata_round_trip_preserves_4xx(status_code: int, error_type: str): + original = OmniClientError("blocked", status_code=status_code, error_type=error_type) + + # Outbound: what the broad `except Exception` handlers serialize. + carried_status, carried_type = client_error_metadata(original) + assert carried_status == status_code + assert carried_type == error_type + assert is_client_error_status(carried_status) + + # Inbound: reconstruction at the consuming end. + rebuilt = client_error_from_metadata("blocked", status_code=carried_status, error_type=carried_type) + assert isinstance(rebuilt, OmniClientError) + assert rebuilt.status_code == status_code + assert rebuilt.error_type == error_type + assert str(rebuilt) == "blocked" + + +@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS) +def test_non_fatal_client_error_preserves_non_400_status(status_code: int, error_type: str): + base = _make_base() + msg = ErrorMessage( + error="client side failure", + status_code=status_code, + error_type=error_type, + ) + + with pytest.raises(OmniClientError) as exc_info: + base._handle_output_message(msg) + + assert exc_info.value.status_code == status_code + assert exc_info.value.error_type == error_type + assert str(exc_info.value) == "client side failure" + + +@pytest.mark.parametrize( + "status_code", + [ + pytest.param(None, id="no-status"), + pytest.param(399, id="399-below-4xx"), + pytest.param(500, id="500-server-error"), + pytest.param(503, id="503-service-unavailable"), + ], +) +def test_non_fatal_non_4xx_status_raises_runtime(status_code: int | None): + base = _make_base() + msg = ErrorMessage( + error="server side failure", + status_code=status_code, + error_type="ShouldBeIgnoredForNon4xx", + ) + + with pytest.raises(RuntimeError) as exc_info: + base._handle_output_message(msg) + + assert not isinstance(exc_info.value, OmniClientError) + assert str(exc_info.value) == "server side failure" + + +def _make_enqueue_client_error( + status_code: int, + error_type: str, + error_text: str, +) -> Callable[[FakeAsyncOmniEngine, dict[str, Any]], None]: + def _enqueue(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None: + engine.output_q.put_nowait( + ErrorMessage( + request_id=msg["request_id"], + stage_id=2, + error=error_text, + status_code=status_code, + error_type=error_type, + ) + ) + + return _enqueue + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS) +async def test_async_omni_propagates_non_400_client_error_status( + monkeypatch: pytest.MonkeyPatch, + status_code: int, + error_type: str, +): + error_text = f"blocked with {status_code}" + engine = FakeAsyncOmniEngine( + stage_metadata=THREE_STAGE_META, + on_add_request=_make_enqueue_client_error(status_code, error_type, error_text), + ) + _patch_engine(monkeypatch, engine) + + app = AsyncOmni("dummy-model") + try: + with pytest.raises(OmniClientError) as exc_info: + async for _ in app.generate(prompt="blocked", request_id="req-1"): + pass + finally: + app.shutdown() + + assert exc_info.value.status_code == status_code + assert exc_info.value.error_type == error_type + assert str(exc_info.value) == error_text + + def test_async_omni_errored_property_alive(): omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index 1365b320d..c2a5ff0a1 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/tests/entrypoints/test_omni_sleep_mode.py @@ -532,7 +532,10 @@ async def test_pure_diffusion_scenario(tp_size: int): try: await engine.sleep(level=1) await engine.wake_up() - async for _ in engine.generate("test", sampling_params=SamplingParams()): + async for _ in engine.generate( + "test", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256), + ): pass logger.info("Pure diffusion OK (TP=%s)", tp_size) finally: diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index f31bf3f20..05593d647 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -135,22 +135,6 @@ def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pyte # ---- _map_device_list unit tests ---- -@pytest.mark.core_model -@pytest.mark.cpu -def test_map_device_list_idempotency(): - """Device IDs already in the visible set are returned as-is (idempotency).""" - result = _map_device_list(0, ["0", "1"], ["0", "1", "2", "3"]) - assert result == ["0", "1"] - - -@pytest.mark.core_model -@pytest.mark.cpu -def test_map_device_list_idempotency_single(): - """Single device ID in visible set passes through.""" - result = _map_device_list(1, ["1"], ["0", "1"]) - assert result == ["1"] - - @pytest.mark.core_model @pytest.mark.cpu def test_map_device_list_index_mapping(): @@ -211,6 +195,18 @@ def test_map_device_list_partial_mapping(): assert result == ["5", "6"] +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_multi_replica_offset_cvd(): + """Logical indices from split_devices_for_replicas must map through an offset + CUDA_VISIBLE_DEVICES, not be treated as physical IDs. + """ + visible = ["2", "3", "4", "5"] + assert _map_device_list(1, ["1"], visible) == ["3"] + assert _map_device_list(1, ["2"], visible) == ["4"] + assert _map_device_list(1, ["2", "3"], visible) == ["4", "5"] + + @pytest.mark.core_model @pytest.mark.cpu def test_map_device_list_raises_on_non_numeric(): diff --git a/tests/helpers/media.py b/tests/helpers/media.py index 5c76a9e1e..e1d995546 100644 --- a/tests/helpers/media.py +++ b/tests/helpers/media.py @@ -1,5 +1,6 @@ """Synthetic media generation and media/text utilities for tests.""" +import atexit import base64 import concurrent.futures import gc @@ -14,7 +15,6 @@ import re import subprocess import tempfile import time -import uuid from contextlib import contextmanager from pathlib import Path from typing import Any @@ -588,15 +588,32 @@ def cosine_similarity_text(text1, text2, n: int = 3): return cosine * length_harmony -def _merge_base64_audio_to_segment(base64_list: list[str]): - from pydub import AudioSegment +class _AudioBuffer: + """Minimal replacement for pydub.AudioSegment used by test helpers.""" - merged = None + def __init__(self, data: np.ndarray, sample_rate: int): + self.data = data + self.sample_rate = sample_rate + + def export(self, buf: io.BytesIO, format: str = "wav"): + sf.write(buf, self.data, self.sample_rate, format=format.upper()) + buf.seek(0) + + +def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer: + from vllm.multimodal.media.audio import AudioMediaIO + + io_ = AudioMediaIO() + chunks: list[np.ndarray] = [] + sample_rate: int | None = None for b64 in base64_list: raw = base64.b64decode(b64.split(",", 1)[-1]) - seg = AudioSegment.from_file(io.BytesIO(raw)) - merged = seg if merged is None else merged + seg - return merged + waveform, sr = io_.load_bytes(raw) + if sample_rate is None: + sample_rate = int(sr) + chunks.append(waveform) + merged = np.concatenate(chunks) if chunks else np.array([], dtype=np.float32) + return _AudioBuffer(merged, sample_rate or 16000) @contextmanager @@ -666,7 +683,10 @@ def convert_audio_file_to_text(output_path: str) -> str: def convert_audio_bytes_to_text(raw_bytes: bytes) -> str: - output_path = f"./test_{uuid.uuid4().hex}.wav" + output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav") + os.close(output_fd) + if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"): + atexit.register(Path(output_path).unlink, missing_ok=True) data, samplerate = sf.read(io.BytesIO(raw_bytes)) sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16") print(f"audio data is saved: {output_path}") diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 807ffb969..e87017df8 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -1473,6 +1473,8 @@ class OpenAIClientHandler: - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server speaks first, e.g. ``/v1/realtime`` rejection path). + - ``ws_skip_types``: optional event ``type`` strings to ignore while waiting for the first matching frame + (e.g. ``["session.created"]`` on ``/v1/realtime``). - ``timeout``: seconds to wait for the first inbound text frame (default ``120``). - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present. """ @@ -1486,6 +1488,7 @@ class OpenAIClientHandler: timeout = float(cfg.get("timeout", 120.0)) uri = self._build_ws_url(path) + skip_types = set(cfg.get("ws_skip_types") or []) connect_kw: dict[str, Any] = {} if "ws_max_size" in cfg: @@ -1497,16 +1500,19 @@ class OpenAIClientHandler: async with websockets.connect(uri, **connect_kw) as ws: for frame in frames: await ws.send(frame) - raw = await asyncio.wait_for(ws.recv(), timeout=timeout) - if not isinstance(raw, str): - raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}") - try: - data = json.loads(raw) - except json.JSONDecodeError as exc: - raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc - if not isinstance(data, dict): - raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}") - return WebSocketJsonResponse(json_body=data) + while True: + raw = await asyncio.wait_for(ws.recv(), timeout=timeout) + if not isinstance(raw, str): + raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}") + try: + data = json.loads(raw) + except json.JSONDecodeError as exc: + raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc + if not isinstance(data, dict): + raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}") + if skip_types and data.get("type") in skip_types: + continue + return WebSocketJsonResponse(json_body=data) resp = asyncio.run(_recv_first_json_object()) _run_ws_expectations_from_request_config(cfg, resp) @@ -1767,7 +1773,16 @@ class OpenAIClientHandler: # Qwen3-TTS custom fields, forwarded via extra_body. extra_body: dict[str, Any] = {} # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params. - for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens", "seed"): + for key in ( + "task_type", + "ref_text", + "ref_audio", + "language", + "max_new_tokens", + "seed", + "instructions", + "speed", + ): if key in request_config: extra_body[key] = request_config[key] diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py index b0afc9592..a2c3142c7 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py @@ -84,6 +84,11 @@ def _make_code2wav_model( token_mel_ratio=2 if with_stride_cfg else 0, ) model.code2wav = _DummyCode2Wav(vocab_size=4, num_samples=num_samples, outputs=outputs) + # Short-circuit the lazy TensorRT estimator swap: these tests exercise the + # forward audio logic, not the TRT path. On a GPU CI runner the swap would + # otherwise run and dereference ``self.model_dir`` (only set in __init__, + # which this fixture bypasses via object.__new__). + model._code2wav_trt_done = True model.source_cache_len = 4 model.speech_window = torch.hamming_window(8, periodic=False) model._stream_audio_cache_by_req = {} diff --git a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py new file mode 100644 index 000000000..7aab6133b --- /dev/null +++ b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig +from vllm_omni.model_executor.models.common.ming.fm import Solver +from vllm_omni.model_executor.models.ming_tts.constants import ( + AGGREGATOR_HIDDEN_SIZE, + HISTORY_PATCH_SIZE, + LATENT_DIM, + LLM_HIDDEN_SIZE, + LLM_VOCAB_SIZE, + PATCH_SIZE, + SAMPLE_RATE, + VAE_PATCH_SIZE, +) +from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss +from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel +from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config +from vllm_omni.model_executor.models.output_templates import OmniOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts] + + +def test_ming_tts_audio_vae_uses_common_config(): + """AudioVAEConfig is shared by Ming dense and Ming flash modules.""" + cfg = AudioVAEConfig(sample_rate=16000, patch_size=-1) + + assert cfg.sample_rate == 16000 + assert cfg.patch_size == -1 + + +def test_ming_tts_cfm_solver_uses_common_implementation(): + """Ming dense imports the shared solver implementation directly.""" + assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm" + + +def test_ming_tts_flowloss_preserves_checkpoint_prefix(): + flowloss = FlowLoss(z_channels=4, llm_cond_dim=8, hidden_size=16, depth=1, num_heads=2) + + assert any(name.startswith("cfm.model.") for name in flowloss.state_dict()) + + +def test_ming_dense_validation_rejects_semantic_audio_vae_config(): + """Dense 0.5B validation rejects semantic AudioVAE configs.""" + cfg = SimpleNamespace( + audio_dummy_token_id=151705, + audio_eos_token_id=151704, + text_eos_token_id=151669, + audio_tokenizer_config=AudioVAEConfig( + sample_rate=SAMPLE_RATE, + patch_size=VAE_PATCH_SIZE, + semantic_module_kwargs={"whisper_encoder": {}}, + enc_kwargs={"latent_dim": LATENT_DIM, "input_dim": 882, "hop_size": 882}, + dec_kwargs={"latent_dim": LATENT_DIM, "output_dim": 882}, + ), + latent_dim=LATENT_DIM, + patch_size=PATCH_SIZE, + history_patch_size=HISTORY_PATCH_SIZE, + llm_hidden_size=LLM_HIDDEN_SIZE, + llm_vocab_size=LLM_VOCAB_SIZE, + sample_rate=SAMPLE_RATE, + vae_patch_size=VAE_PATCH_SIZE, + llm_config={"hidden_size": LLM_HIDDEN_SIZE}, + aggregator_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE}, + ditar_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE}, + latent_chunk_size=1, + latent_left_context=0, + max_decode_steps=1, + stop_head_threshold=0.5, + stop_head_min_steps=0, + ) + + with pytest.raises(ValueError, match="semantic_module_kwargs"): + validate_ming_tts_config(cfg) + + +def test_ming_instruction_parser_preserves_dense_and_flash_defaults(): + """Ming dense and Ming flash keep distinct instruction defaults.""" + serving = object.__new__(OmniOpenAIServingSpeech) + serving.uploaded_speakers = {"uploaded": {}} + + dense_plain = serving._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None)) + assert dense_plain == "calm" + + dense_with_fields = serving._parse_ming_instruction( + SimpleNamespace(instructions="calm", language="Auto", voice="灵小甄") + ) + assert dense_with_fields == {"IP": "灵小甄", "风格": "calm"} + + flash_fields = serving._parse_ming_instruction_fields( + SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄") + ) + assert flash_fields == {"风格": "calm"} + + +def _make_ming_logits_model(vocab_size=8): + model = object.__new__(MingLLMModel) + model.ming_config = SimpleNamespace( + llm_vocab_size=vocab_size, + max_decode_steps=1, + stop_head_min_steps=0, + text_eos_token_id=7, + ) + model._last_text_mode = False + model._last_ming_next_token_ids = None + return model + + +def test_ming_compute_logits_uses_cached_forced_next_token_ids(): + model = _make_ming_logits_model() + model._last_ming_next_token_ids = [2, 5] + + logits = MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace()) + + assert logits.shape == (2, 8) + assert logits[0, 2].item() == 0.0 + assert logits[1, 5].item() == 0.0 + assert torch.isneginf(logits[0, [0, 1, 3, 4, 5, 6, 7]]).all() + assert torch.isneginf(logits[1, [0, 1, 2, 3, 4, 6, 7]]).all() + assert model._last_ming_next_token_ids is None + + +def test_ming_compute_logits_falls_back_to_dummy_token_id(): + model = _make_ming_logits_model() + + logits = MingLLMModel.compute_logits(model, torch.zeros((1, 4)), SimpleNamespace()) + assert logits.shape == (1, 8) + assert logits[0, 7].item() == 0.0 + assert torch.isneginf(logits[0, [0, 1, 2, 3, 4, 5, 6]]).all() + + +def test_ming_forward_non_decode_return_clears_cached_forced_next_token_ids(): + class FakeBackbone: + def __call__(self, **kwargs): + return kwargs["inputs_embeds"] + + model = _make_ming_logits_model() + model.model = FakeBackbone() + model._last_ming_next_token_ids = [2] + + output = MingLLMModel.forward( + model, + input_ids=torch.tensor([1]), + positions=torch.tensor([0]), + inputs_embeds=torch.zeros((1, 4)), + model_intermediate_buffer=[], + ) + + assert isinstance(output, OmniOutput) + assert output.multimodal_outputs is None + assert model._last_ming_next_token_ids is None + + +def test_ming_compute_logits_rejects_forced_token_batch_mismatch(): + model = _make_ming_logits_model() + model._last_ming_next_token_ids = [2] + + with pytest.raises(RuntimeError, match="batch mismatch"): + MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace()) + + +def test_ming_compute_logits_text_mode_delegates_to_backbone(): + class FakeBackbone: + def __init__(self): + self.hidden_states = None + + def compute_logits(self, hidden_states): + self.hidden_states = hidden_states + return torch.ones((hidden_states.shape[0], 3)) + + model = _make_ming_logits_model(vocab_size=3) + model._last_text_mode = True + model.model = FakeBackbone() + hidden_states = torch.zeros((2, 4)) + + logits = MingLLMModel.compute_logits(model, hidden_states, SimpleNamespace()) + + assert torch.equal(model.model.hidden_states, hidden_states) + assert torch.equal(logits, torch.ones((2, 3))) diff --git a/tests/model_executor/models/moss_tts/__init__.py b/tests/model_executor/models/moss_tts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/moss_tts/test_codec_cudagraph.py b/tests/model_executor/models/moss_tts/test_codec_cudagraph.py new file mode 100644 index 000000000..49d6fa411 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_codec_cudagraph.py @@ -0,0 +1,361 @@ +"""Tests for MossTTSCUDAGraphCodecWrapper numerical equivalence. + +Verifies that CUDA Graph-accelerated decoding produces results equivalent +to eager mode, with special attention to the two-argument _decode interface +(codes [NQ, 1, T] + lengths [1]) and the NQ-first input convention. +""" + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import ( + MossAudioTokenizerDecoderOutput, +) +from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( + MossTTSCUDAGraphCodecWrapper, +) + +pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")] + +DEVICE = torch.device("cuda:0") +NUM_QUANTIZERS = 8 +DOWNSAMPLE_RATE = 4 # synthetic; real checkpoint uses 1920 + + +# --------------------------------------------------------------------------- +# Synthetic codec model +# --------------------------------------------------------------------------- + + +class SyntheticCodecModel(nn.Module): + """Minimal stand-in for MossAudioTokenizerModel. + + Exposes the same two-argument _decode(codes, lengths) interface and + returns a MossAudioTokenizerDecoderOutput, mirroring the real model. + + Input: + codes: [NQ, B, T] long — RVQ codes + lengths: [B] long — valid frame counts + Output: + MossAudioTokenizerDecoderOutput(audio=[B, 1, T*upsample], audio_lengths=[B]) + """ + + def __init__( + self, + num_quantizers: int = NUM_QUANTIZERS, + downsample_rate: int = DOWNSAMPLE_RATE, + ): + super().__init__() + self.downsample_rate = downsample_rate + hidden = 32 + # embed: treat NQ quantizers as input channels for Conv1d + self.embed = nn.Conv1d(num_quantizers, hidden, kernel_size=3, padding=1) + self.conv = nn.Conv1d(hidden, hidden, kernel_size=3, padding=1) + self.upsample = nn.ConvTranspose1d(hidden, 1, kernel_size=downsample_rate, stride=downsample_rate) + + def _decode(self, codes: torch.Tensor, lengths: torch.Tensor) -> MossAudioTokenizerDecoderOutput: + """codes: [NQ, B, T], lengths: [B] → audio [B, 1, T*upsample].""" + nq, b, t = codes.shape + # treat NQ as channel dim → [B, NQ, T] for Conv1d (expects [B, C, T]) + x = codes.permute(1, 0, 2).to(dtype=self.embed.weight.dtype) # [B, NQ, T] + x = torch.relu(self.embed(x)) # [B, hidden, T] + x = torch.relu(self.conv(x)) # [B, hidden, T] + audio = self.upsample(x) # [B, 1, T * upsample] + audio_lengths = lengths * self.downsample_rate + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) + + def batch_decode( + self, + codes_list: list[torch.Tensor], + num_quantizers: int | None = None, + ) -> MossAudioTokenizerDecoderOutput: + """Mirrors MossAudioTokenizerModel.batch_decode (fallback path).""" + device = codes_list[0].device + nq = num_quantizers or codes_list[0].shape[0] + max_t = max(c.shape[-1] for c in codes_list) + codes = torch.zeros(nq, len(codes_list), max_t, device=device, dtype=torch.long) + lengths = torch.zeros(len(codes_list), device=device, dtype=torch.long) + for i, c in enumerate(codes_list): + codes[:nq, i, : c.shape[-1]] = c[:nq] + lengths[i] = c.shape[-1] + return self._decode(codes, lengths) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def model(): + torch.manual_seed(0) + return SyntheticCodecModel().to(DEVICE).eval() + + +@pytest.fixture(scope="module") +def wrapper(model): + w = MossTTSCUDAGraphCodecWrapper( + model=model, + capture_sizes=[25, 50, 100], + num_quantizers=NUM_QUANTIZERS, + enabled=True, + ) + w.warmup(DEVICE) + return w + + +def _random_codes(t: int, device: torch.device = DEVICE) -> torch.Tensor: + """Return [NQ, T] long tensor (single-request convention).""" + return torch.randint(0, 64, (NUM_QUANTIZERS, t), dtype=torch.long, device=device) + + +def _eager_decode(model: SyntheticCodecModel, codes_nq_t: torch.Tensor) -> MossAudioTokenizerDecoderOutput: + """Run the model directly in eager mode (reference output).""" + t = codes_nq_t.shape[-1] + codes_nq_1_t = codes_nq_t.unsqueeze(1) # [NQ, 1, T] + lengths = torch.tensor([t], dtype=torch.long, device=codes_nq_t.device) + with torch.no_grad(): + return model._decode(codes_nq_1_t, lengths) + + +# --------------------------------------------------------------------------- +# 1. Exact-size inputs — must be bit-identical +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [25, 50, 100]) +def test_exact_size_bit_identical(model, wrapper, t): + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 2. Padded inputs — output trimmed to actual length +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [10, 30, 47, 73, 99]) +def test_padded_output_shape(model, wrapper, t): + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + expected_len = t * DOWNSAMPLE_RATE + assert out.audio.shape[-1] == expected_len, f"got {out.audio.shape[-1]}, want {expected_len}" + assert out.audio.shape == ref.audio.shape + + +@pytest.mark.parametrize("t", [10, 30, 47, 73, 99]) +def test_padded_interior_positions_close(model, wrapper, t): + """Positions well away from the zero-padding boundary must be numerically close.""" + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + boundary = 2 * DOWNSAMPLE_RATE + if ref.audio.shape[-1] > boundary: + torch.testing.assert_close( + out.audio[..., :-boundary], + ref.audio[..., :-boundary], + atol=1e-5, + rtol=1e-5, + ) + + +# --------------------------------------------------------------------------- +# 3. Fallback to eager (T > all capture sizes) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [101, 150, 200]) +def test_fallback_exact_match(model, wrapper, t): + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 4. Static buffer not aliased across calls +# --------------------------------------------------------------------------- + + +def test_output_not_aliased_after_later_replay(model, wrapper): + """clone() in decode() must prevent later replays from overwriting the result.""" + codes_a = _random_codes(50) + codes_b = _random_codes(50) + with torch.no_grad(): + out_a = wrapper.decode(codes_a) + saved = out_a.audio.clone() + _ = wrapper.decode(codes_b) # replay overwrites static buffer + torch.testing.assert_close(out_a.audio, saved, atol=0, rtol=0) + + +def test_deterministic_across_calls(model, wrapper): + codes = _random_codes(30) + with torch.no_grad(): + out1 = wrapper.decode(codes) + out2 = wrapper.decode(codes) + torch.testing.assert_close(out1.audio, out2.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 5. Disabled wrapper falls back to eager +# --------------------------------------------------------------------------- + + +def test_disabled_falls_back_to_eager(model, wrapper): + codes = _random_codes(30) + ref = _eager_decode(model, codes) + wrapper.enabled = False + with torch.no_grad(): + out = wrapper.decode(codes) + wrapper.enabled = True + torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 6. audio_lengths is consistent with audio shape +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [25, 30, 50, 99, 100]) +def test_audio_lengths_consistent(model, wrapper, t): + codes = _random_codes(t) + with torch.no_grad(): + out = wrapper.decode(codes) + assert out.audio_lengths is not None + assert int(out.audio_lengths[0].item()) == out.audio.shape[-1] + + +# --------------------------------------------------------------------------- +# 7. NQ-first input layout: codes_nq_t is [NQ, T], not [T, NQ] +# --------------------------------------------------------------------------- + + +def test_nq_first_layout_matches_eager(model, wrapper): + """Verify that wrapper correctly interprets [NQ, T] (NQ-first) layout.""" + t = 25 + codes = _random_codes(t) # [NQ, T] + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 8. Non-contiguous input — copy_() must handle strided tensors correctly +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [20, 47]) +def test_noncontiguous_input_matches_contiguous(model, wrapper, t): + """Padded path uses copy_() on a slice; must produce the same result + whether codes_nq_t is contiguous or a non-contiguous view (e.g. a column + slice of a larger buffer, which is the realistic production layout when the + caller slices out one request from a batched tensor).""" + codes_contiguous = _random_codes(t) + # Build a non-contiguous view: embed in a larger [NQ, 2, T] buffer and + # take column 0. The resulting tensor has stride (2*T, 2, 1) → non-contiguous. + buf = torch.zeros(NUM_QUANTIZERS, 2, t, dtype=torch.long, device=DEVICE) + buf[:, 0, :] = codes_contiguous + codes_noncontig = buf[:, 0, :] # [NQ, T], non-contiguous + assert not codes_noncontig.is_contiguous() + + with torch.no_grad(): + out_c = wrapper.decode(codes_contiguous) + out_nc = wrapper.decode(codes_noncontig) + torch.testing.assert_close(out_c.audio, out_nc.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 9. Zero-padding uses token 0 — must be a valid codebook index (no OOB) +# --------------------------------------------------------------------------- + + +def test_zero_padding_no_cuda_error(model, wrapper): + """Padded inputs fill unused frames with code 0. Verify that the zero- + padded region does not corrupt the output shape. + + Note: SyntheticCodecModel uses Conv1d (not nn.Embedding), so it cannot + trigger a real CUDA index-out-of-bounds on large token ids. This test + validates shape correctness only; OOB safety on the real codec requires + an integration test with MossAudioTokenizerModel.""" + t = 7 # well below the 25-frame bucket → large zero-padded region + codes = _random_codes(t) + with torch.no_grad(): + out = wrapper.decode(codes) + torch.accelerator.synchronize() # surface any deferred CUDA errors + assert out.audio.shape[-1] == t * DOWNSAMPLE_RATE + + +# --------------------------------------------------------------------------- +# 10. Shape alignment: graph path and eager fallback return identical shapes +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [25, 30, 101]) +def test_graph_and_eager_shape_identical(model, wrapper, t): + """Graph replay (t<=100) and eager fallback (t>100) must return tensors + with the same number of dimensions and the same batch/channel layout so + downstream consumers can treat both paths uniformly.""" + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + assert out.audio.ndim == ref.audio.ndim, f"ndim mismatch: graph={out.audio.ndim} eager={ref.audio.ndim}" + assert out.audio.shape[:-1] == ref.audio.shape[:-1], ( + f"batch/channel shape mismatch: graph={out.audio.shape} eager={ref.audio.shape}" + ) + + +# --------------------------------------------------------------------------- +# 11. Half-precision weights: graph and eager must agree under fp16/bf16 +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_half_precision_graph_eager_agree(dtype): + """CUDA Graph must produce outputs consistent with eager when model weights + are in half precision. The wrapper itself stores codes as long and lengths + as long; only the model's conv weights change dtype.""" + torch.manual_seed(0) + model_hp = SyntheticCodecModel().to(device=DEVICE, dtype=dtype).eval() + wrapper_hp = MossTTSCUDAGraphCodecWrapper( + model=model_hp, + capture_sizes=[25, 50], + num_quantizers=NUM_QUANTIZERS, + enabled=True, + ) + wrapper_hp.warmup(DEVICE) + + codes = _random_codes(25) + ref = _eager_decode(model_hp, codes) + with torch.no_grad(): + out = wrapper_hp.decode(codes) + # Half-precision accumulation introduces small numerical error; use + # tolerances that match typical fp16/bf16 rounding. + atol = 1e-2 if dtype == torch.float16 else 5e-2 + torch.testing.assert_close(out.audio.float(), ref.audio.float(), atol=atol, rtol=1e-3) + + +# --------------------------------------------------------------------------- +# 12. Empty input (T=0): wrapper must not crash and return zero-length audio +# --------------------------------------------------------------------------- + + +def test_empty_input_t0(model, wrapper): + """T=0 input must be handled gracefully: no exception, audio length == 0. + + The outer MossTTSCodecDecoder guards against empty segments, but the + wrapper should also handle T=0 without crashing so it is safe to call + directly.""" + codes = _random_codes(0) # [NQ, 0] + with torch.no_grad(): + out = wrapper.decode(codes) + assert out.audio.shape[-1] == 0, f"expected empty audio, got shape {out.audio.shape}" diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index 4b2bdf3fe..52ca29905 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -921,6 +921,7 @@ def test_cosyvoice3_full_payload_replace_keys_present() -> None: def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None: """Smoke: ming_flash_omni token-only carries voice metadata.""" from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import ( + thinker2talker, thinker2talker_token_only, ) @@ -938,14 +939,15 @@ def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None: src = [_Wrap("hello world")] prompt = _Prompt({"voice_name": "ZH_FEMALE", "prompt_text": "ref text"}) - out = thinker2talker_token_only(src, prompt=prompt) - assert len(out) == 1 - assert out[0]["prompt_token_ids"] == [0] # talker self-tokenizes; dummy id - info = out[0]["additional_information"] - assert info["text"] == "hello world" - assert info["voice_name"] == "ZH_FEMALE" - assert info["prompt_text"] == "ref text" - assert info["ming_task"] == "omni" + for func in (thinker2talker, thinker2talker_token_only): + out = func(src, prompt=prompt) + assert len(out) == 1 + assert out[0]["prompt_token_ids"] == [0] + info = out[0]["additional_information"] + assert info["text"] == "hello world" + assert info["voice_name"] == "ZH_FEMALE" + assert info["prompt_text"] == "ref text" + assert info["ming_task"] == "omni" def test_qwen2_5_omni_thinker2talker_token_only_smoke() -> None: diff --git a/tests/reasoning/__init__.py b/tests/reasoning/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/reasoning/test_step_audio_reasoning_parser.py b/tests/reasoning/test_step_audio_reasoning_parser.py new file mode 100644 index 000000000..7eb9fa17a --- /dev/null +++ b/tests/reasoning/test_step_audio_reasoning_parser.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Tests for StepAudioReasoningParser. + +Step-Audio supports two representations of thinking markers: + +1. **Special tokens**: ``<|THINK_START|>`` and ``<|THINK_END|>`` + (single-token IDs in vocab). + +2. **Text markers**: ```` and ```` + (multi-token sequences for , e.g. [522, 26865]). + +These tests use standalone mock objects so they can run without +installing the full vLLM/torch stack. +""" + +import sys +from types import ModuleType +from unittest.mock import MagicMock + +import pytest + +# --------------------------------------------------------------------------- +# Bootstrap minimal vLLM stubs so the parser can be imported +# --------------------------------------------------------------------------- + + +def _install_vllm_stubs() -> dict[str, ModuleType | None]: + """Temporarily install minimal stub modules for the vllm imports the + parser needs, so it can be imported without the full vLLM/torch stack. + + Returns a snapshot of the original ``sys.modules`` entries that were + replaced. The caller MUST pass it to ``_restore_vllm_stubs`` once the + parser is imported. These stubs are process-global: if they leak, they + shadow the real vllm modules for every other test in the same worker + (and any forked server subprocess), e.g. replacing + ``vllm.entrypoints.openai.engine.protocol`` with a stub that lacks + ``ErrorResponse`` and breaking unrelated server-startup tests. + """ + stubs: dict[str, ModuleType] = {} + + # DeltaMessage stub + delta_mod = ModuleType("vllm.entrypoints.openai.engine.protocol") + + class DeltaMessage: + __slots__ = ("reasoning", "content") + + def __init__(self, reasoning=None, content=None): + self.reasoning = reasoning + self.content = content + + delta_mod.DeltaMessage = DeltaMessage + stubs["vllm.entrypoints.openai.engine.protocol"] = delta_mod + + # ReasoningParser stub + reasoning_mod = ModuleType("vllm.reasoning") + + class ReasoningParser: + def __init__(self, tokenizer, *args, **kwargs): + self.model_tokenizer = tokenizer + + @property + def vocab(self): + if not self.model_tokenizer: + return {} + return self.model_tokenizer.get_vocab() + + class _ReasoningParserManager: + reasoning_parsers = {} + lazy_parsers = {} + + @classmethod + def register_lazy_module(cls, name, module_path, class_name): + cls.lazy_parsers[name] = (module_path, class_name) + + reasoning_mod.ReasoningParser = ReasoningParser + reasoning_mod.ReasoningParserManager = _ReasoningParserManager + stubs["vllm.reasoning"] = reasoning_mod + stubs["vllm.reasoning.abs_reasoning_parsers"] = reasoning_mod + + # TokenizerLike stub + tokenizers_mod = ModuleType("vllm.tokenizers") + tokenizers_mod.TokenizerLike = object + stubs["vllm.tokenizers"] = tokenizers_mod + + # ChatCompletionRequest / ResponsesRequest stubs + chat_mod = ModuleType("vllm.entrypoints.openai.chat_completion.protocol") + chat_mod.ChatCompletionRequest = MagicMock + stubs["vllm.entrypoints.openai.chat_completion.protocol"] = chat_mod + responses_mod = ModuleType("vllm.entrypoints.openai.responses.protocol") + responses_mod.ResponsesRequest = MagicMock + stubs["vllm.entrypoints.openai.responses.protocol"] = responses_mod + + # Parent packages: only stub the ones not already importable, so we never + # clobber a real (installed) vllm package tree. + for mod_path in [ + "vllm", + "vllm.entrypoints", + "vllm.entrypoints.openai", + "vllm.entrypoints.openai.engine", + "vllm.entrypoints.openai.chat_completion", + "vllm.entrypoints.openai.responses", + ]: + if mod_path not in sys.modules and mod_path not in stubs: + stubs[mod_path] = ModuleType(mod_path) + + saved = {name: sys.modules.get(name) for name in stubs} + sys.modules.update(stubs) + return saved + + +def _restore_vllm_stubs(saved: dict[str, ModuleType | None]) -> None: + """Undo ``_install_vllm_stubs`` so ``sys.modules`` is left untouched.""" + for name, original in saved.items(): + if original is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = original + + +# Install stubs only for the duration of the import. The parser binds these +# stub classes into its own module namespace at import time, so it keeps using +# them at runtime even after sys.modules is restored below. +_saved_modules = _install_vllm_stubs() +try: + from vllm_omni.reasoning.step_audio_reasoning_parser import ( # noqa: E402 + StepAudioReasoningParser, + ) +finally: + _restore_vllm_stubs(_saved_modules) + +# Short aliases for marker constants +TS = StepAudioReasoningParser.THINK_START_TEXT # +TE = StepAudioReasoningParser.THINK_END_TEXT # +TSS = StepAudioReasoningParser.THINK_START_SPECIAL # <|THINK_START|> +TES = StepAudioReasoningParser.THINK_END_SPECIAL # <|THINK_END|> + + +# --------------------------------------------------------------------------- +# Mock tokenizers +# --------------------------------------------------------------------------- + + +class SingleTokenTokenizer: + """Tokenizer where both ```` and ```` are single tokens. + + This exercises the fast (token-ID) path. + Also includes special-token forms ``<|THINK_START|>`` and ``<|THINK_END|>``. + """ + + def __init__(self): + self._vocab = { + TS: 100, # + TE: 101, # + TSS: 151669, # <|THINK_START|> + TES: 151670, # <|THINK_END|> + "Hello": 1, + " world": 2, + "reasoning": 3, + " content": 4, + "\n": 5, + "": 7, + "some text": 8, + } + self._reverse = {v: k for k, v in self._vocab.items()} + + def get_vocab(self): + return dict(self._vocab) + + def decode(self, ids, skip_special_tokens=False): + return "".join(self._reverse.get(i, "") for i in ids) + + def tokenize(self, text): + if not text: + return [] + return [text] + + +class MultiTokenTokenizer: + """Tokenizer where ```` is NOT a single token. + + ```` is in the vocab (ID 100), but ```` is NOT. + Instead, ```` encodes as [""] → [522, 26865]. + + Both special tokens ``<|THINK_START|>`` and ``<|THINK_END|>`` ARE + present as single tokens. + + This exercises the text-based path (with special-token fallback). + """ + + def __init__(self): + self._vocab = { + TS: 100, # - single token + # is intentionally NOT in the vocab + "": 26865, + TSS: 151669, # <|THINK_START|> + TES: 151670, # <|THINK_END|> + "Hello": 1, + " world": 2, + "reasoning": 3, + " content": 4, + "\n": 5, + "some text": 8, + } + self._reverse = {v: k for k, v in self._vocab.items()} + + def get_vocab(self): + return dict(self._vocab) + + def decode(self, ids, skip_special_tokens=False): + return "".join(self._reverse.get(i, "") for i in ids) + + def tokenize(self, text): + if not text: + return [] + return [text] + + +class TextOnlyTokenizer: + """Tokenizer with only text-form markers (no special tokens at all). + + Both ```` and ```` are single tokens. + ``<|THINK_START|>`` and ``<|THINK_END|>`` are NOT in the vocab. + + This tests the pure text-based fallback paths. + """ + + def __init__(self): + self._vocab = { + TS: 100, # + TE: 101, # + "Hello": 1, + " world": 2, + "reasoning": 3, + " content": 4, + "\n": 5, + "": 7, + } + self._reverse = {v: k for k, v in self._vocab.items()} + + def get_vocab(self): + return dict(self._vocab) + + def decode(self, ids, skip_special_tokens=False): + return "".join(self._reverse.get(i, "") for i in ids) + + def tokenize(self, text): + if not text: + return [] + return [text] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _stream_extract(parser, deltas_and_ids): + """Run streaming extraction over a list of (delta_text, delta_ids) tuples. + Returns (reasoning_text, content_text).""" + prev_text = "" + prev_ids = [] + all_r = [] + all_c = [] + + for delta_text, delta_ids in deltas_and_ids: + cur_text = prev_text + delta_text + cur_ids = prev_ids + delta_ids + result = parser.extract_reasoning_streaming(prev_text, cur_text, delta_text, prev_ids, cur_ids, delta_ids) + if result is not None: + if result.reasoning: + all_r.append(result.reasoning) + if result.content: + all_c.append(result.content) + prev_text = cur_text + prev_ids = cur_ids + + return "".join(all_r), "".join(all_c) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def single_token_parser(): + return StepAudioReasoningParser(SingleTokenTokenizer()) + + +@pytest.fixture +def multi_token_parser(): + return StepAudioReasoningParser(MultiTokenTokenizer()) + + +@pytest.fixture +def text_only_parser(): + return StepAudioReasoningParser(TextOnlyTokenizer()) + + +@pytest.fixture +def request_obj(): + return MagicMock(messages=[], model="test-model") + + +# --------------------------------------------------------------------------- +# Initialization tests +# --------------------------------------------------------------------------- + + +class TestInit: + def test_single_token_ids(self, single_token_parser): + """Text-form markers are in vocab as single tokens.""" + assert single_token_parser.think_start_text_id == 100 + assert single_token_parser.think_end_text_id == 101 + + def test_single_token_special_ids(self, single_token_parser): + """Special-token markers are also in vocab.""" + assert single_token_parser.think_start_special_id == 151669 + assert single_token_parser.think_end_special_id == 151670 + + def test_single_token_has_ids(self, single_token_parser): + assert single_token_parser._has_start_token_id is True + assert single_token_parser._has_end_token_id is True + + def test_multi_token_text_ids(self, multi_token_parser): + """ is a single token; is NOT.""" + assert multi_token_parser.think_start_text_id == 100 + assert multi_token_parser.think_end_text_id == -1 + + def test_multi_token_special_ids(self, multi_token_parser): + """Special tokens are still in vocab.""" + assert multi_token_parser.think_start_special_id == 151669 + assert multi_token_parser.think_end_special_id == 151670 + + def test_multi_token_has_ids(self, multi_token_parser): + """Even without as single token, special tokens provide IDs.""" + assert multi_token_parser._has_start_token_id is True + assert multi_token_parser._has_end_token_id is True + + def test_text_only_ids(self, text_only_parser): + """Only text-form markers, no special tokens.""" + assert text_only_parser.think_start_text_id == 100 + assert text_only_parser.think_end_text_id == 101 + assert text_only_parser.think_start_special_id == -1 + assert text_only_parser.think_end_special_id == -1 + assert text_only_parser._has_start_token_id is True + assert text_only_parser._has_end_token_id is True + + def test_missing_tokenizer_raises(self): + with pytest.raises(ValueError, match="model tokenizer must be passed"): + StepAudioReasoningParser(None) + + +# --------------------------------------------------------------------------- +# extract_reasoning (non-streaming) +# --------------------------------------------------------------------------- + + +class TestExtractReasoning: + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_basic_extraction(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = f"some reasoning{TE}Hello world" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "some reasoning" + assert content == "Hello world" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_no_end_token(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = "still thinking..." + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "still thinking..." + assert content is None + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_with_start_token_prefix(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = f"{TS}my reasoning{TE}answer" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "my reasoning" + assert content == "answer" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_strips_newline_after_end(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = f"reasoning{TE}\nHello" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "reasoning" + assert content == "Hello" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_empty_reasoning(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = f"{TE}Hello" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning is None + assert content == "Hello" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + ], + ) + def test_special_end_token(self, request, parser_name, request_obj): + """Test with <|THINK_END|> special token form.""" + parser = request.getfixturevalue(parser_name) + output = f"reasoning{TES}answer" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "reasoning" + assert content == "answer" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + ], + ) + def test_special_start_token(self, request, parser_name, request_obj): + """Test with <|THINK_START|> special token form.""" + parser = request.getfixturevalue(parser_name) + output = f"{TSS}reasoning{TE}answer" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "reasoning" + assert content == "answer" + + +# --------------------------------------------------------------------------- +# is_reasoning_end +# --------------------------------------------------------------------------- + + +class TestIsReasoningEnd: + def test_single_token_text_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end([1, 2, 101]) is True + + def test_single_token_special_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end([1, 2, 151670]) is True + + def test_single_token_no_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end([1, 2, 3]) is False + + def test_multi_token_text_end_decoded(self, multi_token_parser): + """ = [522, 26865] detected via decode.""" + assert multi_token_parser.is_reasoning_end([1, 522, 26865]) is True + + def test_multi_token_special_end(self, multi_token_parser): + """<|THINK_END|> detected as single token ID.""" + assert multi_token_parser.is_reasoning_end([1, 151670]) is True + + def test_multi_token_no_end(self, multi_token_parser): + assert multi_token_parser.is_reasoning_end([1, 2, 3]) is False + + def test_multi_token_partial_marker(self, multi_token_parser): + """Only should NOT count.""" + assert multi_token_parser.is_reasoning_end([1, 522, 3]) is False + + def test_multi_token_empty_input(self, multi_token_parser): + assert multi_token_parser.is_reasoning_end([]) is False + + def test_text_only_end(self, text_only_parser): + assert text_only_parser.is_reasoning_end([1, 2, 101]) is True + + def test_text_only_no_end(self, text_only_parser): + assert text_only_parser.is_reasoning_end([1, 2, 3]) is False + + def test_start_after_end_returns_false(self, single_token_parser): + """Multi-turn prompt: start marker AFTER end marker → reasoning + is still active (the generation prompt's start marker).""" + # IDs: [start=100, end=101, start=100] → last marker is start + assert single_token_parser.is_reasoning_end([100, 101, 100]) is False + + def test_end_is_last_marker_returns_true(self, single_token_parser): + """End marker is the last marker → reasoning has ended.""" + # IDs: [start=100, end=101] → last marker is end + assert single_token_parser.is_reasoning_end([100, 101]) is True + + def test_start_only_no_end(self, single_token_parser): + """Only a start marker, no end marker → reasoning is active.""" + assert single_token_parser.is_reasoning_end([100, 1, 2]) is False + + def test_multi_turn_prompt_start_after_end(self, multi_token_parser): + """Simulated multi-turn prompt with previous assistant response. + + Prompt contains: ... " + " " + ", 151669=<|THINK_START|>] + ids = [151670, 151669] + assert multi_token_parser.is_reasoning_end(ids) is False + + +# --------------------------------------------------------------------------- +# is_reasoning_end_streaming +# --------------------------------------------------------------------------- + + +class TestIsReasoningEndStreaming: + def test_single_token_delta_has_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end_streaming([1, 2], [101]) is True + + def test_single_token_delta_no_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False + + def test_multi_token_delta_has_both_parts(self, multi_token_parser): + """Full in delta → [522, 26865].""" + assert multi_token_parser.is_reasoning_end_streaming([1], [522, 26865]) is True + + def test_multi_token_marker_spans_boundary(self, multi_token_parser): + """ in delta.""" + assert multi_token_parser.is_reasoning_end_streaming([1, 522], [26865]) is True + + def test_multi_token_no_end(self, multi_token_parser): + assert multi_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False + + +# --------------------------------------------------------------------------- +# extract_reasoning_streaming +# --------------------------------------------------------------------------- + + +class TestExtractReasoningStreaming: + """Test streaming extraction for both single-token and multi-token modes.""" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_simple_streaming(self, request_, parser_name): + parser = request_.getfixturevalue(parser_name) + # Simulate: reasoning final + if parser_name == "multi_token_parser": + deltas = [TS, "reasoning", TE, " content"] + token_ids_per_delta = [ + [100], # + [3], # reasoning + [522, 26865], # + [4], # content + ] + else: + deltas = [TS, "reasoning", TE, " content"] + token_ids_per_delta = [ + [100], # + [3], # reasoning + [101], # + [4], # content + ] + + r, c = _stream_extract(parser, list(zip(deltas, token_ids_per_delta))) + assert r == "reasoning" + assert c == " content" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_no_end_token_streaming(self, request_, parser_name): + parser = request_.getfixturevalue(parser_name) + deltas_and_ids = [ + (TS, [100]), + ("thinking", [3]), + (" more", [2]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "thinking more" + assert c == "" + + def test_multi_token_end_split_across_chunks(self, multi_token_parser): + """The key test: is split across streaming chunks. + + Previous chunk ends with "reasoning\\nfinal answer". This is the real-world scenario when + ```` is tokenized as two separate tokens. + """ + deltas_and_ids = [ + ("reasoning", [3]), + ("\nHello", [26865, 1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "Hello" + + def test_multi_token_start_token_in_delta(self, multi_token_parser): + """Test that start token is skipped in streaming output.""" + deltas_and_ids = [ + (TS, [100]), + ("reasoning", [3]), + (TE, [522, 26865]), + ("Hello", [1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "Hello" + + def test_multi_token_three_way_split(self, multi_token_parser): + """ split across 3 chunks: .""" + deltas_and_ids = [ + ("some reasoning", [3]), + ("", [8]), + ("\nHello", [1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "some reasoning" + assert c == "Hello" + + def test_special_end_token_in_stream(self, multi_token_parser): + """<|THINK_END|> special token in streaming (decoded as text).""" + deltas_and_ids = [ + ("reasoning", [3]), + (TES, [151670]), + ("answer", [1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "answer" + + def test_no_start_token_in_stream(self, multi_token_parser): + """Model starts directly with reasoning (no prefix).""" + deltas_and_ids = [ + ("reasoning", [3]), + (TE, [522, 26865]), + ("answer", [1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "answer" + + def test_empty_reasoning(self, single_token_parser): + """Start immediately followed by end → empty reasoning.""" + deltas_and_ids = [ + (TS, [100]), + (TE, [101]), + ("Hello", [1]), + ] + r, c = _stream_extract(single_token_parser, deltas_and_ids) + assert r in ("", None) + assert c == "Hello" + + def test_end_marker_no_content_after(self, multi_token_parser): + """End marker with no content following it.""" + deltas_and_ids = [ + ("reasoning", [3]), + (TE, [522, 26865]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "" + + def test_delta_is_just_start_token(self, multi_token_parser): + """Delta consisting only of the start token should be skipped.""" + deltas_and_ids = [ + (TS, [100]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "" + assert c == "" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + ], + ) + def test_only_first_newline_stripped_after_end(self, request_, parser_name): + """Only the first newline after should be stripped.""" + parser = request_.getfixturevalue(parser_name) + if parser_name == "multi_token_parser": + deltas_and_ids = [ + ("reasoning", [3]), + (TE, [522, 26865]), + ("\n\nHello", [5, 1]), + ] + else: + deltas_and_ids = [ + ("reasoning", [3]), + (TE, [101]), + ("\n\nHello", [5, 1]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reasoning" + assert c == "\nHello" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + ], + ) + def test_content_after_end_multiple_deltas(self, request_, parser_name): + """Multiple content deltas after end marker.""" + parser = request_.getfixturevalue(parser_name) + if parser_name == "multi_token_parser": + deltas_and_ids = [ + (TS, [100]), + ("reasoning", [3]), + (TE, [522, 26865]), + ("Hello", [1]), + (" world", [2]), + ] + else: + deltas_and_ids = [ + (TS, [100]), + ("reasoning", [3]), + (TE, [101]), + ("Hello", [1]), + (" world", [2]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reasoning" + assert c == "Hello world" + + def test_pending_buffer_resolved_on_non_marker(self, multi_token_parser): + """When buffered text turns out not to be a marker prefix, it is + emitted as reasoning on the next call.""" + # "reason<" is buffered because "<" could start a marker. + # Next delta ">" doesn't form a marker, so the buffer is + # flushed as reasoning. + parser = multi_token_parser + deltas_and_ids = [ + (TS, [100]), + ("reason", [3]), + ("<", [8]), # "<" is buffered (prefix of markers) + ("> more", [8]), # ">" doesn't complete any marker + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reason> more" + assert c == "" + + def test_multi_token_end_with_reasoning_before_buffer(self, multi_token_parser): + """Reasoning text before a buffered prefix should be emitted.""" + # "reasoning<" → "reasoning" emitted, "<" buffered + # "\nanswer" → buffer resolves as end marker + parser = multi_token_parser + deltas_and_ids = [ + ("reasoning", [3]), + ("<", [8]), + (f"/{TE[2:]}\nanswer", [522, 26865, 1]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reasoning" + assert c == "answer" + + def test_multi_token_five_way_split(self, multi_token_parser): + """ split across 5 tiny chunks.""" + # Split: each chunk is one character of + chunks = list(TE) + deltas_and_ids = [ + ("reasoning", [3]), + ] + for ch in chunks: + deltas_and_ids.append((ch, [8])) + deltas_and_ids.append(("\nHello", [1])) + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "Hello" + + def test_single_token_parser_with_pending(self, single_token_parser): + """Single-token parser: ambiguous '<' prefix buffered then resolved.""" + # Even with single-token markers, the text-based path buffers + # '<' because it could be a prefix of a marker. + parser = single_token_parser + deltas_and_ids = [ + (TS, [100]), + ("reason", [3]), + ("<", [8]), + ("> not a marker", [8]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reason> not a marker" + assert c == "" + + +# --------------------------------------------------------------------------- + + +class TestCountReasoningTokens: + def test_single_token_basic(self, single_token_parser): + # IDs: [100=, 3=reasoning, 101=, 4=content] + assert single_token_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1 + + def test_single_token_no_end(self, single_token_parser): + # All tokens after start are reasoning + assert single_token_parser.count_reasoning_tokens([100, 3, 2]) == 2 + + def test_multi_token_basic(self, multi_token_parser): + # encodes as [522, 26865] + # Full sequence: reasoning content + # IDs: [100, 3, 522, 26865, 4] + result = multi_token_parser.count_reasoning_tokens([100, 3, 522, 26865, 4]) + assert result > 0 + + def test_multi_token_no_end(self, multi_token_parser): + # No end token → all tokens inside are counted. + result = multi_token_parser.count_reasoning_tokens([100, 3, 2]) + assert result > 0 + + def test_text_only_basic(self, text_only_parser): + assert text_only_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1 + + def test_text_only_no_end(self, text_only_parser): + assert text_only_parser.count_reasoning_tokens([100, 3, 2]) == 2 + + +# --------------------------------------------------------------------------- +# extract_content_ids +# --------------------------------------------------------------------------- + + +class TestExtractContentIds: + def test_single_token_basic(self, single_token_parser): + # [1, 101=, 2, 3] → content starts after ID 101 + result = single_token_parser.extract_content_ids([1, 101, 2, 3]) + assert result == [2, 3] + + def test_single_token_no_end(self, single_token_parser): + result = single_token_parser.extract_content_ids([1, 2, 3]) + assert result == [] + + def test_multi_token_basic(self, multi_token_parser): + # = [522, 26865] + # [3, 522, 26865, 4] → content after + result = multi_token_parser.extract_content_ids([3, 522, 26865, 4]) + assert 4 in result + + def test_multi_token_no_end(self, multi_token_parser): + result = multi_token_parser.extract_content_ids([1, 2, 3]) + assert result == [] + + def test_text_only_basic(self, text_only_parser): + result = text_only_parser.extract_content_ids([1, 101, 2, 3]) + assert result == [2, 3] + + def test_text_only_no_end(self, text_only_parser): + result = text_only_parser.extract_content_ids([1, 2, 3]) + assert result == [] + + +# --------------------------------------------------------------------------- +# Fixture request workaround (pytest "request" fixture shadowed by parametrize) +# --------------------------------------------------------------------------- + + +@pytest.fixture +def request_(request): + """Proxy for the pytest ``request`` fixture to avoid name collisions + with the ``request_obj`` fixture used in extract_reasoning tests.""" + return request diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index e541d3628..e3ebea88f 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -17,7 +17,6 @@ from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.config.stage_config import ( _PIPELINE_REGISTRY, DeployConfig, - ModelPipeline, PipelineConfig, StageConfig, StageConfigFactory, @@ -351,93 +350,6 @@ class TestStageConfig: assert "parallel_config" not in omega_config.engine_args -class TestModelPipeline: - """Tests for ModelPipeline class.""" - - def test_valid_linear_dag(self): - """Test validation of a valid linear DAG.""" - stages = [ - StageConfig(stage_id=0, model_stage="thinker", input_sources=[]), - StageConfig(stage_id=1, model_stage="talker", input_sources=[0]), - StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]), - ] - pipeline = ModelPipeline(model_type="test", stages=stages) - errors = pipeline.validate_pipeline() - assert errors == [], f"Unexpected errors: {errors}" - - def test_valid_branching_dag(self): - """Test validation of a valid branching DAG.""" - stages = [ - StageConfig(stage_id=0, model_stage="input", input_sources=[]), - StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]), - StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]), - ] - pipeline = ModelPipeline(model_type="test", stages=stages) - errors = pipeline.validate_pipeline() - assert errors == [], f"Unexpected errors: {errors}" - - def test_missing_entry_point(self): - """Test that missing entry point is detected.""" - stages = [ - StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]), - StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]), - ] - pipeline = ModelPipeline(model_type="test", stages=stages) - errors = pipeline.validate_pipeline() - assert any("entry point" in e.lower() for e in errors) - - def test_missing_dependency(self): - """Test that missing stage reference is detected.""" - stages = [ - StageConfig(stage_id=0, model_stage="input", input_sources=[]), - StageConfig(stage_id=1, model_stage="output", input_sources=[99]), # Invalid - ] - pipeline = ModelPipeline(model_type="test", stages=stages) - errors = pipeline.validate_pipeline() - assert any("non-existent" in e.lower() for e in errors) - - def test_duplicate_stage_ids(self): - """Test that duplicate stage IDs are detected.""" - stages = [ - StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]), - StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]), # Duplicate - ] - pipeline = ModelPipeline(model_type="test", stages=stages) - errors = pipeline.validate_pipeline() - assert any("duplicate" in e.lower() for e in errors) - - def test_self_reference(self): - """Test that self-references are detected.""" - stages = [ - StageConfig(stage_id=0, model_stage="entry", input_sources=[]), - StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]), # Self - ] - pipeline = ModelPipeline(model_type="test", stages=stages) - errors = pipeline.validate_pipeline() - assert any("itself" in e.lower() for e in errors) - - def test_get_stage_by_id(self): - """Test getting stage by ID.""" - stages = [ - StageConfig(stage_id=0, model_stage="thinker", input_sources=[]), - StageConfig(stage_id=1, model_stage="talker", input_sources=[0]), - ] - pipeline = ModelPipeline(model_type="test", stages=stages) - - stage = pipeline.get_stage(1) - assert stage is not None - assert stage.model_stage == "talker" - - missing = pipeline.get_stage(99) - assert missing is None - - def test_empty_pipeline(self): - """Test validation of empty pipeline.""" - pipeline = ModelPipeline(model_type="test", stages=[]) - errors = pipeline.validate_pipeline() - assert any("no stages" in e.lower() for e in errors) - - class TestStageConfigFactory: """Tests for StageConfigFactory class.""" @@ -598,304 +510,6 @@ class TestStageResolutionHelpers: assert overridden == [] -class TestPipelineYamlParsing: - """Tests for pipeline YAML file parsing (@ZJY0516).""" - - def test_parse_qwen3_omni_moe_yaml(self, tmp_path): - """Test parsing the qwen3_omni_moe pipeline YAML.""" - yaml_content = """\ -model_type: qwen3_omni_moe -async_chunk: false - -stages: - - stage_id: 0 - model_stage: thinker - stage_type: llm - input_sources: [] - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - hf_config_name: thinker_config - final_output: true - final_output_type: text - is_comprehension: true - - - stage_id: 1 - model_stage: talker - stage_type: llm - input_sources: [0] - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - hf_config_name: talker_config - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker - - - stage_id: 2 - model_stage: code2wav - stage_type: llm - input_sources: [1] - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - hf_config_name: thinker_config - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav - final_output: true - final_output_type: audio -""" - yaml_file = tmp_path / "qwen3_omni_moe.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe") - - assert pipeline.model_type == "qwen3_omni_moe" - assert len(pipeline.stages) == 3 - - # Stage 0: thinker - s0 = pipeline.stages[0] - assert s0.stage_id == 0 - assert s0.model_stage == "thinker" - assert s0.stage_type == StageType.LLM - assert s0.input_sources == [] - assert s0.worker_type == "ar" - assert s0.final_output is True - assert s0.final_output_type == "text" - assert s0.is_comprehension is True - - # Stage 1: talker - s1 = pipeline.stages[1] - assert s1.stage_id == 1 - assert s1.input_sources == [0] - assert s1.custom_process_input_func == ( - "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker" - ) - assert s1.final_output is False - - # Stage 2: code2wav - s2 = pipeline.stages[2] - assert s2.stage_id == 2 - assert s2.input_sources == [1] - assert s2.worker_type == "generation" - assert s2.final_output_type == "audio" - - def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path): - """Test backward compatibility with engine_input_source field.""" - yaml_content = """\ -model_type: legacy_model - -stages: - - stage_id: 0 - model_stage: entry - stage_type: llm - - stage_id: 1 - model_stage: downstream - stage_type: llm - engine_input_source: [0] -""" - yaml_file = tmp_path / "legacy.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model") - assert pipeline.stages[1].input_sources == [0] - - def test_parse_yaml_with_connectors_and_edges(self, tmp_path): - """Test parsing pipeline with optional connectors and edges.""" - yaml_content = """\ -model_type: test_model - -stages: - - stage_id: 0 - model_stage: entry - stage_type: llm - input_sources: [] - -connectors: - type: ray - -edges: - - from: 0 - to: 1 -""" - yaml_file = tmp_path / "with_connectors.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model") - assert pipeline.connectors == {"type": "ray"} - assert pipeline.edges == [{"from": 0, "to": 1}] - - def test_parsed_pipeline_passes_validation(self, tmp_path): - """Test that a well-formed YAML produces a valid pipeline.""" - yaml_content = """\ -model_type: valid_model - -stages: - - stage_id: 0 - model_stage: entry - stage_type: llm - input_sources: [] - final_output: true - final_output_type: text - - stage_id: 1 - model_stage: next - stage_type: llm - input_sources: [0] -""" - yaml_file = tmp_path / "valid.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model") - errors = pipeline.validate_pipeline() - assert errors == [], f"Unexpected validation errors: {errors}" - - def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path): - """Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs.""" - yaml_content = """\ -model_type: legacy_batch -stages: - - stage_id: 0 - model_stage: entry - stage_type: llm - input_sources: [] - runtime: - devices: "0" - max_batch_size: 16 - engine_args: - model_arch: SomeModel -""" - yaml_file = tmp_path / "legacy_batch.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch") - s0 = pipeline.stages[0] - assert "max_batch_size" not in s0.yaml_runtime - assert s0.yaml_engine_args.get("max_num_seqs") == 16 - - def test_parse_diffusion_stage_type(self, tmp_path): - """Test parsing a diffusion stage type from YAML.""" - yaml_content = """\ -model_type: diff_model - -stages: - - stage_id: 0 - model_stage: dit - stage_type: diffusion - input_sources: [] - final_output: true - final_output_type: image -""" - yaml_file = tmp_path / "diffusion.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model") - assert pipeline.stages[0].stage_type == StageType.DIFFUSION - - def test_parse_mimo_audio_yaml(self, tmp_path): - """Test parsing the MiMo Audio pipeline YAML.""" - yaml_content = """\ -model_type: mimo_audio - -stages: - - stage_id: 0 - model_stage: fused_thinker_talker - stage_type: llm - input_sources: [] - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - is_comprehension: true - final_output: true - final_output_type: text - - - stage_id: 1 - model_stage: code2wav - stage_type: llm - input_sources: [0] - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav - final_output: true - final_output_type: audio -""" - yaml_file = tmp_path / "mimo_audio.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio") - - assert pipeline.model_type == "mimo_audio" - assert len(pipeline.stages) == 2 - - # Stage 0: fused_thinker_talker - s0 = pipeline.stages[0] - assert s0.stage_id == 0 - assert s0.model_stage == "fused_thinker_talker" - assert s0.input_sources == [] - assert s0.worker_type == "ar" - assert s0.is_comprehension is True - assert s0.final_output is True - assert s0.final_output_type == "text" - - # Stage 1: code2wav - s1 = pipeline.stages[1] - assert s1.stage_id == 1 - assert s1.input_sources == [0] - assert s1.worker_type == "generation" - assert s1.custom_process_input_func == ( - "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav" - ) - assert s1.final_output is True - assert s1.final_output_type == "audio" - - # Pipeline validation - errors = pipeline.validate_pipeline() - assert errors == [], f"Unexpected errors: {errors}" - - -class TestAsyncChunk: - """Tests for async_chunk pipeline flag.""" - - def test_model_pipeline_async_chunk_default(self): - """Test ModelPipeline defaults async_chunk to False.""" - stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])] - pipeline = ModelPipeline(model_type="test", stages=stages) - assert pipeline.async_chunk is False - - def test_model_pipeline_async_chunk_set(self): - """Test ModelPipeline with async_chunk=True.""" - stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])] - pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True) - assert pipeline.async_chunk is True - - def test_parse_async_chunk_from_yaml(self, tmp_path): - """Test that async_chunk is parsed from pipeline YAML.""" - yaml_content = """\ -model_type: test_async -async_chunk: true - -stages: - - stage_id: 0 - model_stage: entry - stage_type: llm - input_sources: [] -""" - yaml_file = tmp_path / "async.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async") - assert pipeline.async_chunk is True - - def test_parse_missing_async_chunk_defaults_false(self, tmp_path): - """Test that missing async_chunk defaults to False.""" - yaml_content = """\ -model_type: test_no_async - -stages: - - stage_id: 0 - model_stage: entry - stage_type: llm - input_sources: [] -""" - yaml_file = tmp_path / "no_async.yaml" - yaml_file.write_text(yaml_content) - - pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async") - assert pipeline.async_chunk is False - - class TestPipelineDiscovery: """Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``).""" diff --git a/tests/test_pydub_removal.py b/tests/test_pydub_removal.py new file mode 100644 index 000000000..ae65bdab2 --- /dev/null +++ b/tests/test_pydub_removal.py @@ -0,0 +1,182 @@ +"""Functional tests for pydub removal. + +Verifies that the new av/numpy/torchaudio replacements produce +identical results to the old pydub-based code for all three paths: + 1. WAV PCM → duration & frame count (pure arithmetic) + 2. Non-WAV (mp3/ogg) → duration & frame count (av.open) + 3. TTS WER PCM export: resample to 24000Hz mono s16 +""" + +import io +import wave + +import av +import numpy as np +import soundfile as sf + + +def _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=1.0): + """Generate a WAV file with a sine wave, return (wav_bytes, pcm_bytes, params).""" + n_samples = int(sample_rate * duration_sec) + t = np.linspace(0, duration_sec, n_samples, endpoint=False) + samples = (np.sin(2 * np.pi * 440 * t) * 32767 * 0.5).astype(np.int16) + if channels == 2: + samples = np.column_stack([samples, samples]).flatten() + + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(channels) + wf.setsampwidth(sample_width) + wf.setframerate(sample_rate) + wf.writeframes(samples.tobytes()) + wav_bytes = buf.getvalue() + + # Extract raw PCM (skip WAV header) + with wave.open(io.BytesIO(wav_bytes), "rb") as wr: + pcm_bytes = wr.readframes(wr.getnframes()) + params = (wr.getnchannels(), wr.getsampwidth(), wr.getframerate()) + + return wav_bytes, pcm_bytes, params + + +def _make_ogg_bytes(sample_rate=16000, channels=1, duration_sec=1.0): + """Generate an OGG Vorbis file with a sine wave, return bytes.""" + n_samples = int(sample_rate * duration_sec) + t = np.linspace(0, duration_sec, n_samples, endpoint=False) + samples = (np.sin(2 * np.pi * 440 * t) * 0.5).astype(np.float32) + buf = io.BytesIO() + sf.write(buf, samples, sample_rate, format="OGG", subtype="VORBIS") + return buf.getvalue() + + +def test_wav_path_duration_and_frames(): + """Path 1: WAV PCM buffer → duration & frames via arithmetic.""" + for sr in [16000, 24000, 44100]: + for ch in [1, 2]: + for dur in [0.5, 1.0, 2.5]: + _, pcm_bytes, (channels, sample_width, frame_rate) = _make_wav_bytes( + sample_rate=sr, channels=ch, duration_sec=dur + ) + + # New code logic (from patch.py) + frame_width = sample_width * channels + audio_frames = len(pcm_bytes) // frame_width + audio_duration_sec = audio_frames / frame_rate + + assert abs(audio_duration_sec - dur) < 0.001, f"sr={sr} ch={ch} dur={dur}: got {audio_duration_sec}" + expected_frames = int(sr * dur) + assert audio_frames == expected_frames, ( + f"sr={sr} ch={ch} dur={dur}: frames {audio_frames} != {expected_frames}" + ) + + print("[PASS] test_wav_path_duration_and_frames") + + +def test_non_wav_path_duration_av(): + """Path 2: non-WAV audio → duration & frames via av.""" + for dur in [0.5, 1.0, 2.0]: + ogg_bytes = _make_ogg_bytes(sample_rate=16000, duration_sec=dur) + + container = av.open(io.BytesIO(ogg_bytes)) + stream = container.streams.audio[0] + audio_frames = sum(f.samples for f in container.decode(stream)) + audio_duration_sec = audio_frames / stream.rate if stream.rate else 0.0 + container.close() + + # OGG encoding may add padding, allow small tolerance + assert abs(audio_duration_sec - dur) < 0.05, f"dur={dur}: got {audio_duration_sec}" + + print("[PASS] test_non_wav_path_duration_av") + + +def test_tts_pcm_export_wav_path(): + """Path 3a: WAV PCM → resample to 24000Hz mono s16.""" + import torch + import torchaudio + + _, pcm_bytes, (channels, sample_width, frame_rate) = _make_wav_bytes( + sample_rate=16000, channels=2, sample_width=2, duration_sec=1.0 + ) + + # New code logic + pcm_array = np.frombuffer(pcm_bytes, dtype=np.int16) + if channels > 1: + pcm_array = pcm_array.reshape(-1, channels).mean(axis=1).astype(np.int16) + if frame_rate != 24000: + t = torch.from_numpy(pcm_array.astype(np.float32)).unsqueeze(0) + t = torchaudio.functional.resample(t, frame_rate, 24000) + pcm_array = t.squeeze(0).to(torch.int16).numpy() + + result_bytes = pcm_array.tobytes() + # Should be mono, 24000Hz, 2 bytes per sample + expected_samples = int(24000 * 1.0) + actual_samples = len(result_bytes) // 2 + assert abs(actual_samples - expected_samples) < 10, f"expected ~{expected_samples} samples, got {actual_samples}" + assert len(pcm_array.shape) == 1, "should be mono (1D array)" + + print("[PASS] test_tts_pcm_export_wav_path") + + +def test_tts_pcm_export_non_wav_path(): + """Path 3b: non-WAV → resample to 24000Hz mono s16 via av.AudioResampler.""" + ogg_bytes = _make_ogg_bytes(sample_rate=16000, channels=1, duration_sec=1.0) + + container = av.open(io.BytesIO(ogg_bytes)) + resampler = av.AudioResampler(format="s16", layout="mono", rate=24000) + frames = [] + for frame in container.decode(audio=0): + for rf in resampler.resample(frame): + frames.append(rf.to_ndarray().flatten()) + container.close() + + result = np.concatenate(frames) + result_bytes = result.tobytes() + expected_samples = int(24000 * 1.0) + actual_samples = len(result_bytes) // 2 + assert abs(actual_samples - expected_samples) < 500, f"expected ~{expected_samples} samples, got {actual_samples}" + + print("[PASS] test_tts_pcm_export_non_wav_path") + + +def test_media_merge_and_export(): + """media.py: _merge_base64_audio_to_segment decode + concat + export WAV.""" + import base64 as b64mod + import sys + + sys.path.insert(0, ".") + from tests.helpers.media import _merge_base64_audio_to_segment + + # Create two WAV chunks as base64 + chunks_b64 = [] + for dur in [0.5, 0.3]: + wav_bytes, _, _ = _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=dur) + chunks_b64.append(b64mod.b64encode(wav_bytes).decode()) + + merged = _merge_base64_audio_to_segment(chunks_b64) + + # Test export to WAV + buf = io.BytesIO() + merged.export(buf, format="wav") + wav_out = buf.getvalue() + assert len(wav_out) > 44, f"WAV output too small: {len(wav_out)} bytes" + + # Verify merged duration (~0.8s) + out_data, out_sr = sf.read(io.BytesIO(wav_out)) + actual_dur = len(out_data) / out_sr + assert abs(actual_dur - 0.8) < 0.05, f"expected ~0.8s merged duration, got {actual_dur}" + + print(f"[PASS] test_media_merge_and_export (duration={actual_dur:.3f}s, wav_size={len(wav_out)} bytes)") + + +if __name__ == "__main__": + test_wav_path_duration_and_frames() + test_non_wav_path_duration_av() + import importlib.util + + if importlib.util.find_spec("torch") and importlib.util.find_spec("torchaudio"): + test_tts_pcm_export_wav_path() + else: + print("[SKIP] test_tts_pcm_export_wav_path (torch/torchaudio not available)") + test_tts_pcm_export_non_wav_path() + test_media_merge_and_export() + print("\nAll tests passed.") diff --git a/tests/unit/higgs_audio_v3/__init__.py b/tests/unit/higgs_audio_v3/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py new file mode 100644 index 000000000..066d0c6a0 --- /dev/null +++ b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py @@ -0,0 +1,1364 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic unit tests for higgs-audio v3. + +These tests verify AC-1 (config), AC-3 (prompt), AC-4 (fused modules), +AC-5 (delay pattern), AC-7 (stage processor), and AC-10 (registry) +without requiring the actual checkpoint or GPU. +""" + +import asyncio +import time +from collections import OrderedDict + +import pytest +import torch + +# ---- AC-1: Configuration ---- + + +class TestHiggsAudioV3Config: + def test_default_config_loads(self): + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + + config = HiggsAudioV3Config() + assert config.num_codebooks == 8 + assert config.codebook_size == 1026 + assert config.audio_stream_bos_id == 1024 + assert config.audio_stream_eos_id == 1025 + assert config.sample_rate == 24000 + assert config.frame_rate == 25 + assert config.num_real_codes == 1024 + assert config.tie_modality_embeddings is True + + def test_custom_audio_encoder_config(self): + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + + config = HiggsAudioV3Config( + audio_encoder_config={ + "encoder_type": "discrete", + "num_codebooks": 4, + "vocab_size": 512, + "tie_word_embeddings": False, + } + ) + assert config.num_codebooks == 4 + assert config.codebook_size == 512 + assert config.tie_modality_embeddings is False + + def test_invalid_num_codebooks_rejected(self): + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + + with pytest.raises(ValueError, match="num_codebooks must be > 0"): + HiggsAudioV3Config(audio_encoder_config={"num_codebooks": 0, "vocab_size": 1026}) + + def test_negative_num_codebooks_rejected(self): + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + + with pytest.raises(ValueError, match="num_codebooks must be > 0"): + HiggsAudioV3Config(audio_encoder_config={"num_codebooks": -1, "vocab_size": 1026}) + + def test_special_token_ids_initially_none(self): + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + + config = HiggsAudioV3Config() + assert config.tts_token_id is None + assert config.text_token_id is None + assert config.audio_continuation_id is None + + def test_hidden_size_from_text_config(self): + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + + config = HiggsAudioV3Config(text_config={"model_type": "qwen3", "hidden_size": 2560}) + assert config.hidden_size == 2560 + assert config.audio_hidden_size == 2560 + + +# ---- AC-4: Fused Multi-Codebook Modules ---- + + +class TestFusedModules: + def test_embedding_shape(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsFusedMultiTextEmbedding, + ) + + embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=256) + assert embed.weight.shape == (8 * 1026, 256) + + codes = torch.randint(0, 1024, (5, 8)) + out = embed(codes) + assert out.shape == (5, 256) + + def test_embedding_offset_indexing(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsFusedMultiTextEmbedding, + ) + + embed = HiggsFusedMultiTextEmbedding(num_codebooks=2, vocab_size=4, hidden_size=3) + # Set weights to identity-like pattern so we can verify offsets + with torch.no_grad(): + embed.weight.zero_() + # Codebook 0, vocab [0..3] -> row 0..3 + # Codebook 1, vocab [0..3] -> row 4..7 + for i in range(8): + embed.weight[i, i % 3] = float(i) + + # codes[0] = [1, 2] -> embed[1] + embed[4+2=6] + codes = torch.tensor([[1, 2]]) + out = embed(codes) + expected = embed.weight[1] + embed.weight[6] + assert torch.allclose(out[0], expected) + + def test_head_shape(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsFusedMultiTextHead, + ) + + head = HiggsFusedMultiTextHead(num_codebooks=8, vocab_size=1026, hidden_size=256) + assert head.weight.shape == (8 * 1026, 256) + + hidden = torch.randn(3, 256) + out = head.generate(hidden) + assert out.shape == (3, 8, 1026) + + def test_tying(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsFusedMultiTextEmbedding, + HiggsFusedMultiTextHead, + ) + + embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=64) + head = HiggsFusedMultiTextHead(num_codebooks=8, vocab_size=1026, hidden_size=64) + head.weight = embed.weight + assert head.weight is embed.weight + # Modification to one should reflect in the other + with torch.no_grad(): + embed.weight[0, 0] = 42.0 + assert head.weight[0, 0] == 42.0 + + def test_codes_out_of_range(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsFusedMultiTextEmbedding, + ) + + embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16) + # Code 1026 + offset for codebook 7 = 1026 + 7*1026 = 8208 + # But weight only has 8*1026 = 8208 rows, so index 8208 is out of bounds + codes = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 1026]]) # Last cb: 1026 is OOB + with pytest.raises(IndexError): + embed(codes) + + +# ---- AC-5: Delay Pattern Behavior ---- + + +class TestDelayPatternBehavior: + """Test the delay pattern masking logic extracted from the talker.""" + + def test_boc_eoc_ids(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + BOC_ID, + EOC_ID, + ) + + assert BOC_ID == 1024 + assert EOC_ID == 1025 + + def test_delay_phase_boc_masking(self): + """During delay phase, codebooks beyond delay_count are forced to BOC.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID + + num_codebooks = 8 + # Simulate delay phase at step 3 (delay_count=2, so CBs 3-7 should be BOC) + codes = torch.randint(0, 1024, (num_codebooks,)) + delay_count = 2 + next_cb = delay_count + 1 + if next_cb < num_codebooks: + codes[next_cb:] = BOC_ID + # CBs 0-2 should have original codes, CBs 3-7 should be BOC + assert all(codes[i] != BOC_ID or i >= 3 for i in range(num_codebooks)) + assert all(codes[i] == BOC_ID for i in range(3, num_codebooks)) + + def test_cb0_eoc_triggers_rampdown(self): + """EOC on codebook 0 starts ramp-down; EOC on other codebooks does not.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + EOC_ID, + ) + + num_codebooks = 8 + codes = torch.randint(0, 1024, (num_codebooks,)) + + # Simulate: cb0 emits EOC + codes[0] = EOC_ID + assert int(codes[0].item()) == EOC_ID + # Ramp-down should start with countdown = N-2 = 6 + eoc_countdown = num_codebooks - 2 + assert eoc_countdown == 6 + + # If cb3 emits EOC but cb0 doesn't, no ramp-down + codes2 = torch.randint(0, 1024, (num_codebooks,)) + codes2[3] = EOC_ID + assert int(codes2[0].item()) != EOC_ID # cb0 is not EOC + + def test_rampdown_termination(self): + """After N-2 steps of ramp-down, generation_done becomes True.""" + num_codebooks = 8 + eoc_countdown = num_codebooks - 2 # = 6 + generation_done = False + for _ in range(6): + eoc_countdown -= 1 + if eoc_countdown <= 0: + generation_done = True + assert generation_done is True + + def test_all_boc_seed(self): + """First audio step should seed last_codes with all-BOC.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID + + seeded = torch.full((8,), BOC_ID, dtype=torch.long) + assert seeded.shape == (8,) + assert all(seeded[i] == BOC_ID for i in range(8)) + + +# ---- AC-5: Real Sampler Method Tests ---- + + +class TestSamplerMethods: + """Test the actual sampler and batched delay masking methods.""" + + def _make_minimal_talker(self): + """Create a minimal talker-like object with sampler/masking methods.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + class FakeTalker: + num_codebooks = 8 + codebook_size = 1026 + + t = FakeTalker() + t._sample_audio_codes = mod.HiggsAudioV3TalkerForConditionalGeneration._sample_audio_codes.__get__(t) + return t + + def _make_batched_sampler_talker(self, num_rows=4): + """Create a fake talker with GPU-resident sampler state helpers.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + class FakeTalker: + num_codebooks = 8 + codebook_size = 1026 + + t = FakeTalker() + t._decode_last_codes = torch.arange(num_rows * t.num_codebooks, dtype=torch.long).view( + num_rows, t.num_codebooks + ) + t._decode_has_codes = torch.tensor([True, False, True, True])[:num_rows].clone() + t._decode_delay_count = torch.tensor([0, 2, 7, 4], dtype=torch.int32)[:num_rows].clone() + t._decode_eoc_countdown = torch.tensor([-1, -1, 3, 1], dtype=torch.int32)[:num_rows].clone() + t._decode_generation_done = torch.tensor([False, False, False, True])[:num_rows].clone() + t._decode_active_audio_count = 0 + t._codebook_index_cache = {} + t._row_index_cache = {} + t._boc_frame_cache = {} + t._last_audio_codes_buffer = None + t._last_audio_host_staging = None + t._last_audio_gpu_staging = None + t._last_audio_staging_event = None + t._audio_staging_events = [] + t._audio_staging_event_cursor = 0 + t._last_audio_codes = None + t._last_audio_code_valid = [] + t._postprocess_cursor = 0 + t._postprocess_audio_rows = 0 + t._postprocess_audio_active_rows = 0 + + cls = mod.HiggsAudioV3TalkerForConditionalGeneration + for name in ( + "_sample_audio_codes", + "_ensure_decode_state_capacity", + "_get_row_indices", + "_get_codebook_indices", + "_get_boc_frames", + "_get_audio_codes_buffer", + "_get_audio_gpu_staging_buffer", + "_get_audio_host_staging_buffer", + "_apply_delay_pattern_masking_batched", + "_update_delay_state_batched", + "_prefill_row_mask", + "_audio_seed_mask_from_step_input", + ): + setattr(t, name, getattr(cls, name).__get__(t)) + t._device_cache_key = cls._device_cache_key + return t + + def test_sample_respects_mask(self): + """Tokens masked to -inf must never be sampled.""" + t = self._make_batched_sampler_talker(num_rows=1) + logits = torch.full((10, 1026), float("-inf")) + # Only token 500 is allowed for each row + logits[:, 500] = 0.0 + result = t._sample_audio_codes(logits) + assert result.shape == (10,) + assert (result == 500).all(), f"Expected all 500, got {result.tolist()}" + + def test_sample_all_masked_falls_back_to_argmax(self): + """All-masked row should fall back to argmax (least-negative logit).""" + t = self._make_batched_sampler_talker(num_rows=1) + logits = torch.full((2, 1026), float("-inf")) + # Row 0: all masked + # Row 1: only token 42 allowed + logits[1, 42] = 0.0 + result = t._sample_audio_codes(logits) + assert result.shape == (2,) + assert result[1].item() == 42 + + def test_delay_masking_forces_boc_during_delay(self): + """During delay phase, codebooks beyond delay_count must have only BOC allowed.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID + + t = self._make_batched_sampler_talker(num_rows=1) + # Simulate delay_count=2 for batch row 0 + t._ensure_decode_state_capacity(1, torch.device("cpu")) + t._decode_delay_count[0] = 2 + t._decode_eoc_countdown[0] = -1 + cb_logits = torch.zeros(1, 8, 1026) # [1 audio row, 8 codebooks, 1026 vocab] + t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long)) + # CBs 3-7 should have only BOC allowed (everything else -inf) + for q in range(3, 8): + row = cb_logits[0, q] + assert row[BOC_ID].item() == 0.0 # BOC kept + # All non-BOC should be -inf + non_boc = torch.cat([row[:BOC_ID], row[BOC_ID + 1 :]]) + assert (non_boc == float("-inf")).all(), f"CB{q} has non-inf non-BOC values" + + def test_delay_masking_disallows_boc_for_active_codebooks(self): + """Active codebooks during delay should have BOC disallowed.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID + + t = self._make_batched_sampler_talker(num_rows=1) + t._ensure_decode_state_capacity(1, torch.device("cpu")) + t._decode_delay_count[0] = 3 + t._decode_eoc_countdown[0] = -1 + cb_logits = torch.zeros(1, 8, 1026) + t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long)) + # CBs 0-3 are active: BOC should be -inf + for q in range(4): + assert cb_logits[0, q, BOC_ID].item() == float("-inf") + + def test_delay_masking_only_cb0_allows_eoc(self): + """Only codebook 0 should allow EOC during normal generation.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID + + t = self._make_batched_sampler_talker(num_rows=1) + t._ensure_decode_state_capacity(1, torch.device("cpu")) + t._decode_delay_count[0] = 8 + t._decode_eoc_countdown[0] = -1 + cb_logits = torch.zeros(1, 8, 1026) + t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long)) + # CB0 should keep EOC + assert cb_logits[0, 0, EOC_ID].item() != float("-inf") + # CB1-7 should have EOC masked + for q in range(1, 8): + assert cb_logits[0, q, EOC_ID].item() == float("-inf"), f"CB{q} EOC not masked" + + def test_rampdown_masking_locks_to_eoc(self): + """During ramp-down, locked codebooks should only allow EOC.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + BOC_ID, + EOC_ID, + ) + + t = self._make_batched_sampler_talker(num_rows=1) + # Ramp-down with 4 remaining delays: lock CBs 0-3 to EOC + t._ensure_decode_state_capacity(1, torch.device("cpu")) + t._decode_delay_count[0] = 8 + t._decode_eoc_countdown[0] = 4 + cb_logits = torch.zeros(1, 8, 1026) + t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long)) + # CBs 0-3 locked: only EOC allowed + for q in range(4): + row = cb_logits[0, q] + assert row[EOC_ID].item() == 0.0 + non_eoc = torch.cat([row[:EOC_ID], row[EOC_ID + 1 :]]) + assert (non_eoc == float("-inf")).all(), f"Locked CB{q} has non-inf non-EOC" + # CBs 4-7 active: BOC and EOC disallowed + for q in range(4, 8): + assert cb_logits[0, q, BOC_ID].item() == float("-inf") + assert cb_logits[0, q, EOC_ID].item() == float("-inf") + + def test_batched_delay_masking_all_rows_matches_sparse_path(self): + """The all-row fast path must preserve sparse path masking semantics.""" + torch.manual_seed(0) + sparse = self._make_batched_sampler_talker() + all_rows = self._make_batched_sampler_talker() + rows = torch.arange(4, dtype=torch.long) + logits_sparse = torch.randn(4, 8, 1026) + logits_all_rows = logits_sparse.clone() + + sparse._apply_delay_pattern_masking_batched(logits_sparse, rows, all_rows=False) + all_rows._apply_delay_pattern_masking_batched(logits_all_rows, rows, all_rows=True) + + assert torch.equal(logits_all_rows, logits_sparse) + + def test_delay_state_update_all_rows_matches_sparse_path(self): + """The all-row fast path must update state and staging like index_copy_.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + EOC_ID, + ) + + sparse = self._make_batched_sampler_talker() + all_rows = self._make_batched_sampler_talker() + rows = torch.arange(4, dtype=torch.long) + codes = torch.tensor( + [ + [10, 11, 12, 13, 14, 15, 16, 17], + [20, 21, 22, 23, 24, 25, 26, 27], + [30, 31, EOC_ID, 33, 34, 35, 36, 37], + [40, 41, 42, 43, 44, 45, 46, 47], + ], + dtype=torch.long, + ) + code_row_mask = torch.tensor([True, False, True, True]) + + sparse._update_delay_state_batched( + codes.clone(), rows, 4, torch.device("cpu"), code_row_mask=code_row_mask, all_rows=False + ) + all_rows._update_delay_state_batched( + codes.clone(), rows, 4, torch.device("cpu"), code_row_mask=code_row_mask, all_rows=True + ) + + assert torch.equal(all_rows._decode_delay_count, sparse._decode_delay_count) + assert torch.equal(all_rows._decode_eoc_countdown, sparse._decode_eoc_countdown) + assert torch.equal(all_rows._decode_generation_done, sparse._decode_generation_done) + assert torch.equal(all_rows._decode_last_codes, sparse._decode_last_codes) + assert torch.equal(all_rows._decode_has_codes, sparse._decode_has_codes) + assert all_rows._last_audio_codes is None + assert sparse._last_audio_codes is None + assert torch.equal(all_rows._last_audio_host_staging, sparse._last_audio_host_staging) + + def test_direct_audio_sampler_does_not_override_done_tokens(self): + """Done rows are represented directly in GPU sampled tokens.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + t = self._make_batched_sampler_talker(num_rows=2) + t._resolve_token_ids = lambda: None + t._audio_continuation_id = 99999 + t._eos_token_id = 151671 + t._last_logits_hidden = torch.zeros(2, 16) + t._last_step_input_ids = torch.tensor([99999, 99999]) + t._last_step_query_start_loc = None + t._decode_has_codes = torch.tensor([True, True]) + t._decode_generation_done = torch.tensor([False, True]) + t._decode_delay_count = torch.zeros(2, dtype=torch.long) + t._decode_eoc_countdown = torch.full((2,), -1, dtype=torch.long) + t._fast_audio_direct_rows = 2 + t._last_audio_done_flags = None + t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: None + t._audio_codebook_logits_from_rows = lambda hidden, rows, all_rows=False: torch.zeros(2, 8, 1026) + t._apply_delay_pattern_masking_batched = lambda cb_logits, audio_rows, all_rows=False: None + t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long) + t._update_delay_state_batched = lambda *args, **kwargs: None + t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t) + + sampler_output = t.sample(torch.zeros(2, 200000), sampling_metadata=object()) + + assert not getattr( + mod.HiggsAudioV3TalkerForConditionalGeneration, "supports_sampled_token_ids_cpu_override", False + ) + assert sampler_output.sampled_token_ids.tolist() == [[99999], [151671]] + + def test_direct_audio_sampler_uses_gpu_tokens_for_active_rows(self): + """Direct audio sampling leaves scheduler-visible tokens to the GPU tensor.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + t = self._make_batched_sampler_talker(num_rows=2) + t._resolve_token_ids = lambda: None + t._audio_continuation_id = 99999 + t._eos_token_id = 151671 + t._last_logits_hidden = torch.zeros(2, 16) + t._last_step_input_ids = torch.tensor([99999, 99999]) + t._last_step_query_start_loc = None + t._decode_has_codes = torch.tensor([True, True]) + t._decode_generation_done = torch.tensor([False, False]) + t._decode_delay_count = torch.zeros(2, dtype=torch.long) + t._decode_eoc_countdown = torch.full((2,), -1, dtype=torch.long) + t._fast_audio_direct_rows = 2 + t._last_audio_done_flags = [0, 0] + t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: None + t._audio_codebook_logits_from_rows = lambda hidden, rows, all_rows=False: torch.zeros(2, 8, 1026) + t._apply_delay_pattern_masking_batched = lambda cb_logits, audio_rows, all_rows=False: None + t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long) + t._update_delay_state_batched = lambda *args, **kwargs: None + t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t) + + sampler_output = t.sample(torch.zeros(2, 200000), sampling_metadata=object()) + + assert sampler_output.sampled_token_ids.tolist() == [[99999], [99999]] + + def test_terminal_rampdown_frame_is_emitted_before_done(self): + """The final ramp-down frame must be emitted, not replaced by a done marker.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID + + t = self._make_batched_sampler_talker(num_rows=1) + t._decode_delay_count[0] = 8 + t._decode_eoc_countdown[0] = 1 + t._decode_generation_done[0] = False + t._decode_has_codes[0] = True + codes = torch.tensor([[10, 11, 12, 13, 14, 15, 16, 17]], dtype=torch.long) + + t._update_delay_state_batched( + codes, + torch.tensor([0], dtype=torch.long), + 1, + torch.device("cpu"), + code_row_mask=torch.tensor([True]), + all_rows=True, + ) + + staged = t._last_audio_host_staging[0] + assert staged[t.num_codebooks].item() == 1 # valid + assert staged[t.num_codebooks + 1].item() == 1 # done + assert not (staged[: t.num_codebooks] == -1).any() + assert staged[: t.num_codebooks - 1].eq(EOC_ID).all() + assert t._decode_generation_done[0].item() is True + assert t._decode_has_codes[0].item() is False + + def test_decode_state_growth_preserves_inflight_rows(self): + """Growing decode buffers must copy old rows and initialize new rows inactive.""" + t = self._make_batched_sampler_talker(num_rows=2) + old_codes = t._decode_last_codes.clone() + old_has = t._decode_has_codes.clone() + old_delay = t._decode_delay_count.clone() + old_rem = t._decode_eoc_countdown.clone() + old_done = t._decode_generation_done.clone() + + t._ensure_decode_state_capacity(5, torch.device("cpu")) + + assert t._decode_last_codes.shape[0] >= 5 + assert torch.equal(t._decode_last_codes[:2], old_codes) + assert torch.equal(t._decode_has_codes[:2], old_has) + assert torch.equal(t._decode_delay_count[:2], old_delay) + assert torch.equal(t._decode_eoc_countdown[:2], old_rem) + assert torch.equal(t._decode_generation_done[:2], old_done) + assert not t._decode_has_codes[2:5].any() + assert not t._decode_generation_done[2:5].any() + assert t._decode_eoc_countdown[2:5].eq(-1).all() + + def test_mixed_batch_prefill_mask_targets_request_rows(self): + """Mixed prefill/decode must reset only prefill request rows.""" + t = self._make_batched_sampler_talker(num_rows=3) + t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5]) + + mask = t._prefill_row_mask(3, torch.device("cpu")) + + assert torch.equal(mask, torch.tensor([False, True, False])) + + def test_mixed_batch_seed_mask_maps_token_tails_to_request_rows(self): + """Seed detection should return one flag per request, not per token.""" + t = self._make_batched_sampler_talker(num_rows=3) + t._audio_continuation_id = 99999 + t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5]) + t._last_step_input_ids = torch.tensor([99999, 11, 12, 13, 99999]) + + mask = t._audio_seed_mask_from_step_input(3, torch.device("cpu")) + + assert torch.equal(mask, torch.tensor([True, False, True])) + + def test_mixed_batch_seed_mask_allows_prefill_tail_audio_token(self): + """A prefill span ending in <|audio|> starts audio-code generation.""" + t = self._make_batched_sampler_talker(num_rows=3) + t._audio_continuation_id = 99999 + t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5]) + t._last_step_input_ids = torch.tensor([99999, 11, 12, 99999, 99999]) + + mask = t._audio_seed_mask_from_step_input(3, torch.device("cpu")) + + assert torch.equal(mask, torch.tensor([True, True, True])) + + +# ---- AC-6: Feedback Method Tests ---- + + +class TestFeedbackMethods: + def test_postprocess_emits_audio_codes(self): + """postprocess() should return codes from _last_audio_codes.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + class FakeTalker: + _last_audio_codes = torch.tensor([[100, 200, 300, 400, 500, 600, 700, 800]]) + _last_audio_code_valid = [True] + _last_audio_host_staging = None + _last_audio_staging_event = None + _postprocess_cursor = 0 + + t = FakeTalker() + t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t) + t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t) + result = t.postprocess(torch.zeros(1, 64)) + assert "codes" in result + assert "audio" in result["codes"] + assert result["codes"]["audio"].shape == (1, 8) + + def test_postprocess_skips_negative_rows(self): + """postprocess() should skip rows with -1 (no audio).""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + class FakeTalker: + _last_audio_codes = torch.tensor([[-1, -1, -1, -1, -1, -1, -1, -1]]) + _last_audio_code_valid = [False] + _last_audio_host_staging = None + _last_audio_staging_event = None + _postprocess_cursor = 0 + + t = FakeTalker() + t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t) + t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t) + result = t.postprocess(torch.zeros(1, 64)) + assert result == {} + + def test_postprocess_advances_cursor(self): + """postprocess() should advance cursor by 1 per call.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + class FakeTalker: + _last_audio_codes = torch.tensor( + [[100, 200, 300, 400, 500, 600, 700, 800], [-1, -1, -1, -1, -1, -1, -1, -1]] + ) + _last_audio_code_valid = [True, False] + _last_audio_host_staging = None + _last_audio_staging_event = None + _postprocess_cursor = 0 + + t = FakeTalker() + t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t) + t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t) + r1 = t.postprocess(torch.zeros(1, 64)) + assert "codes" in r1 + r2 = t.postprocess(torch.zeros(1, 64)) + assert r2 == {} # Second row is all -1 + + +# ---- AC-6: Audio Feedback Method Tests ---- + + +class TestAudioFeedback: + """Test _apply_audio_feedback with minimal fake talker.""" + + def _make_feedback_talker(self): + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16) + torch.nn.init.ones_(embed.weight) + + class FakeTalker: + num_codebooks = 8 + codebook_size = 1026 + _audio_continuation_id = 99999 # fake audio token + _last_step_query_start_loc = None + _decode_active_audio_count = 0 + _decode_last_codes = torch.zeros(4, 8, dtype=torch.long) + _decode_has_codes = torch.zeros(4, dtype=torch.bool) + _use_external_decode_cudagraph = False + multimodal_embedding = embed + model = type("M", (), {"embed_tokens": lambda self, ids: torch.zeros(ids.shape[0], 16)})() + + t = FakeTalker() + t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t) + t._is_single_token_decode_step = ( + mod.HiggsAudioV3TalkerForConditionalGeneration._is_single_token_decode_step.__get__(t) + ) + t._step_request_count = mod.HiggsAudioV3TalkerForConditionalGeneration._step_request_count.__get__(t) + t._decode_request_token_positions = ( + mod.HiggsAudioV3TalkerForConditionalGeneration._decode_request_token_positions.__get__(t) + ) + t._ensure_decode_state_capacity = lambda min_bs, device: None + return t + + def test_text_positions_unchanged(self): + """Non-audio positions should not be modified.""" + t = self._make_feedback_talker() + # No audio state → all positions unchanged + input_ids = torch.tensor([1, 2, 3, 4]) + hidden = torch.randn(4, 16) + result = t._apply_audio_feedback(hidden, input_ids) + assert torch.equal(result, hidden) + + def test_audio_position_replaced_with_state(self): + """Audio position with state should have its embedding replaced.""" + t = self._make_feedback_talker() + audio_id = t._audio_continuation_id + t._decode_last_codes[1] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[1] = True + t._decode_active_audio_count = 1 + t._ensure_decode_state_capacity = lambda min_bs, device: None + t._last_step_query_start_loc = torch.tensor([0, 1, 2, 3]) + input_ids = torch.tensor([1, audio_id, 3]) + hidden = torch.zeros(3, 16) + result = t._apply_audio_feedback(hidden, input_ids) + # Position 1 (audio) should be non-zero (embedding of all-0 codes) + assert result[1].abs().sum() > 0 + # Positions 0 and 2 should remain zero + assert result[0].abs().sum() == 0 + assert result[2].abs().sum() == 0 + + def test_no_state_no_replacement(self): + """Audio position without state should not be replaced.""" + t = self._make_feedback_talker() + audio_id = t._audio_continuation_id + # No state for row 0 + input_ids = torch.tensor([audio_id]) + hidden = torch.zeros(1, 16) + result = t._apply_audio_feedback(hidden, input_ids) + assert torch.equal(result, hidden) + + def test_prefill_span_does_not_receive_audio_feedback(self): + """A prefill span can have numel == hidden rows but must not be treated as decode.""" + t = self._make_feedback_talker() + t._last_step_input_ids = torch.tensor([1, 2, 3, 4]) + t._last_step_query_start_loc = torch.tensor([0, 4]) + t._decode_last_codes[0] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[0] = True + t._decode_active_audio_count = 1 + hidden = torch.zeros(4, 16) + + result = t._apply_audio_feedback(hidden, t._last_step_input_ids) + + assert torch.equal(result, hidden) + + def test_multi_request_decode_receives_audio_feedback(self): + """Concurrent decode has numel > 1 but each request span is exactly one token.""" + t = self._make_feedback_talker() + t._last_step_input_ids = torch.tensor([10, 11, 12, 13]) + t._last_step_query_start_loc = torch.tensor([0, 1, 2, 3, 4]) + t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[2] = True + t._decode_active_audio_count = 1 + hidden = torch.zeros(4, 16) + + result = t._apply_audio_feedback(hidden, t._last_step_input_ids) + + assert result[2].abs().sum() > 0 + assert result[0].abs().sum() == 0 + assert result[1].abs().sum() == 0 + assert result[3].abs().sum() == 0 + + def test_external_decode_graph_capture_keeps_feedback_ops_without_query_start_loc(self): + """FULL_DECODE capture has no request qsl but still needs dense decode feedback ops.""" + t = self._make_feedback_talker() + t._use_external_decode_cudagraph = True + t._last_step_input_ids = torch.tensor([10, 11, 12, 13]) + t._last_step_query_start_loc = None + t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[2] = True + hidden = torch.zeros(4, 16) + + result = t._apply_audio_feedback(hidden, t._last_step_input_ids) + + assert result[2].abs().sum() > 0 + assert result[0].abs().sum() == 0 + assert result[1].abs().sum() == 0 + assert result[3].abs().sum() == 0 + + def test_mixed_prefill_decode_receives_feedback_only_on_decode_rows(self): + """Mixed batches should skip prefill spans but keep audio feedback for decode spans.""" + t = self._make_feedback_talker() + t._last_step_input_ids = torch.tensor([99999, 21, 22, 23, 99999]) + t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5]) + t._decode_last_codes[0] = torch.zeros(8, dtype=torch.long) + t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[0] = True + t._decode_has_codes[2] = True + t._decode_active_audio_count = 2 + hidden = torch.zeros(5, 16) + + result = t._apply_audio_feedback(hidden, t._last_step_input_ids) + + assert result[0].abs().sum() > 0 + assert result[4].abs().sum() > 0 + assert result[1].abs().sum() == 0 + assert result[2].abs().sum() == 0 + assert result[3].abs().sum() == 0 + + +# ---- AC-8: Codec Strictness ---- + + +class TestCodecStrictness: + def test_bundled_missing_quantizer_key_raises(self): + """Bundled codec load must fail when a quantizer codebook key is missing.""" + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( + HiggsAudioV3Code2Wav, + ) + + config = HiggsAudioV3Config() + c2w = HiggsAudioV3Code2Wav(config=config) + + # State with only 7 quantizers (missing quantizer.quantizers.7.*) + codec_state = {} + for i in range(7): + codec_state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64) + codec_state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64) + codec_state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024) + codec_state["fc2.weight"] = torch.randn(256, 1024) + codec_state["fc2.bias"] = torch.randn(256) + + with pytest.raises(KeyError, match="quantizer 7"): + c2w._load_from_bundled_state(codec_state, device=torch.device("cpu")) + + def test_bundled_missing_fc2_raises(self): + """Bundled codec load must fail when fc2 keys are missing.""" + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( + HiggsAudioV3Code2Wav, + ) + + config = HiggsAudioV3Config() + c2w = HiggsAudioV3Code2Wav(config=config) + + codec_state = {} + for i in range(8): + codec_state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64) + codec_state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64) + codec_state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024) + # No fc2 keys + + with pytest.raises(KeyError, match="fc2"): + c2w._load_from_bundled_state(codec_state, device=torch.device("cpu")) + + def test_bundled_unknown_decoder_key_rejected(self, monkeypatch): + """Unknown decoder-side key should be rejected as unconsumed.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( + HiggsAudioV3Code2Wav, + ) + + class FakeDAC(torch.nn.Module): + def __init__(self): + super().__init__() + self.dummy = torch.nn.Parameter(torch.zeros(1)) + + monkeypatch.setattr(c2w_mod, "build_boson_dac_decoder", lambda device: FakeDAC().to(device)) + monkeypatch.setattr(c2w_mod, "build_higgs_audio_acoustic_decoder", lambda cfg, device: FakeDAC().to(device)) + + config = HiggsAudioV3Config() + c2w = HiggsAudioV3Code2Wav(config=config) + + codec_state = _build_full_codec_state() + codec_state["unknown_decoder_module.weight"] = torch.randn(10) + + with pytest.raises(RuntimeError, match="unexpected decoder-side keys"): + c2w._load_from_bundled_state(codec_state, device=torch.device("cpu")) + + def test_bundled_encoder_side_keys_accepted(self, monkeypatch): + """Known encoder-side keys should not trigger unconsumed-key rejection.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod + from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, + ) + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( + HiggsAudioV3Code2Wav, + ) + + class FakeDAC(torch.nn.Module): + def __init__(self): + super().__init__() + self.dummy = torch.nn.Parameter(torch.zeros(1)) + + monkeypatch.setattr(c2w_mod, "build_boson_dac_decoder", lambda device: FakeDAC().to(device)) + monkeypatch.setattr(c2w_mod, "build_higgs_audio_acoustic_decoder", lambda cfg, device: FakeDAC().to(device)) + + config = HiggsAudioV3Config() + c2w = HiggsAudioV3Code2Wav(config=config) + + codec_state = _build_full_codec_state() + codec_state["acoustic_encoder.block.0.conv_t1.weight"] = torch.randn(10) + codec_state["semantic_model.encoder.layers.0.weight"] = torch.randn(10) + codec_state["fc.weight"] = torch.randn(10) + + # Should NOT raise — encoder-side keys are allowed + c2w._load_from_bundled_state(codec_state, device=torch.device("cpu")) + assert c2w._loaded + + +def _build_minimal_codec_state() -> dict[str, torch.Tensor]: + """Build a minimal codec state dict (quantizer + fc2, no DAC).""" + state: dict[str, torch.Tensor] = {} + for i in range(8): + state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64) + state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64) + state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024) + state["fc2.weight"] = torch.randn(256, 1024) + state["fc2.bias"] = torch.randn(256) + return state + + +def _build_full_codec_state() -> dict[str, torch.Tensor]: + """Build a full codec state dict including acoustic_decoder keys. + + The acoustic_decoder keys here match what a FakeDAC module expects + (just 'dummy'). Real DAC has many more keys. + """ + state = _build_minimal_codec_state() + state["acoustic_decoder.dummy"] = torch.zeros(1) + return state + + +# ---- AC-2: Loader Required-Key Validation ---- + + +class TestLoaderRequiredKeys: + """Test exact required-key validation in load_weights.""" + + def test_build_required_keys_count(self): + """Required key set should have 1 text_emb + 36*11 layers + 1 norm + 1 modality = 399.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration, + ) + + keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(36) + assert len(keys) == 1 + 36 * 11 + 1 + 1 # text + layers + norm + modality = 399 + + def test_build_required_keys_contains_norm(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration, + ) + + keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(2) + assert "body.norm.weight" in keys + + def test_build_required_keys_contains_layer_subkeys(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration, + ) + + keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(2) + assert "body.layers.0.self_attn.k_proj.weight" in keys + assert "body.layers.1.mlp.gate_proj.weight" in keys + + def test_build_required_keys_contains_modality_embedding(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration, + ) + + keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1) + assert "tied.embedding.modality_embeddings.0.embedding.weight" in keys + + def test_build_required_keys_contains_text_embedding(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration, + ) + + keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1) + assert "tied.embedding.text_embedding.weight" in keys + + +class TestLoadWeightsEnforcement: + """Test load_weights() itself with synthetic key sets via a minimal fake talker.""" + + def _make_fake_talker_and_keys(self): + """Build a minimal talker-like object and the full 1-layer key set.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + # Minimal fake talker with just enough to run load_weights + talker = object.__new__(mod.HiggsAudioV3TalkerForConditionalGeneration) + torch.nn.Module.__init__(talker) + + # Attach minimal attributes load_weights needs + talker.multimodal_embedding = mod.HiggsFusedMultiTextEmbedding(8, 1026, 16) + talker.modality_head = mod.HiggsFusedMultiTextHead(8, 1026, 16) + talker.tie_modality = True + talker.modality_head.weight = talker.multimodal_embedding.weight + + # Fake backbone config with 1 layer + class FakeConfig: + num_hidden_layers = 1 + tie_word_embeddings = True + + talker._backbone_config = FakeConfig() + + # Minimal model and lm_head so _BackboneWrapper can be constructed + talker.model = torch.nn.Module() + talker.lm_head = torch.nn.Module() + + # Fake vllm_config for resolve_special_tokens (no-op) + class FakeModelConfig: + model = None + + class FakeVllmConfig: + model_config = FakeModelConfig() + + talker.vllm_config = FakeVllmConfig() + talker.config = type("C", (), {"resolve_special_tokens": lambda self, p: None})() + talker._resolved_tokens = True + talker._audio_continuation_id = None + talker._eos_token_id = None + + # Build the full required key set for 1 layer + required = mod.HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1) + + # Build synthetic weights: tiny tensors for every required key + weights = {} + for key in required: + if "modality_embeddings.0.embedding" in key: + weights[key] = torch.randn(8 * 1026, 16) + else: + weights[key] = torch.randn(4) # tiny placeholder + + # Monkeypatch _BackboneWrapper.load_weights to accept and return keys + def fake_backbone_load(self_wrapper, ws): + consumed = set() + for name, _ in ws: + consumed.add(name) + return consumed + + talker._fake_backbone_load = fake_backbone_load + return talker, weights, mod, fake_backbone_load + + def test_load_weights_succeeds_with_all_keys(self, monkeypatch): + """load_weights succeeds when all required keys are present.""" + talker, weights, mod, fake_load = self._make_fake_talker_and_keys() + monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load) + result = talker.load_weights(iter(weights.items())) + assert isinstance(result, set) + assert "multimodal_embedding.weight" in result + + def test_load_weights_fails_missing_norm(self, monkeypatch): + """load_weights raises when body.norm.weight is missing.""" + talker, weights, mod, fake_load = self._make_fake_talker_and_keys() + monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load) + del weights["body.norm.weight"] + with pytest.raises(RuntimeError, match="body.norm.weight"): + talker.load_weights(iter(weights.items())) + + def test_load_weights_fails_missing_layer_subkey(self, monkeypatch): + """load_weights raises when a specific layer subkey is missing.""" + talker, weights, mod, fake_load = self._make_fake_talker_and_keys() + monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load) + del weights["body.layers.0.self_attn.k_proj.weight"] + with pytest.raises(RuntimeError, match="body.layers.0.self_attn.k_proj.weight"): + talker.load_weights(iter(weights.items())) + + def test_load_weights_fails_missing_modality_embedding(self, monkeypatch): + """load_weights raises when modality embedding is missing.""" + talker, weights, mod, fake_load = self._make_fake_talker_and_keys() + monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load) + del weights["tied.embedding.modality_embeddings.0.embedding.weight"] + with pytest.raises(RuntimeError, match="modality_embeddings"): + talker.load_weights(iter(weights.items())) + + def test_load_weights_fails_unknown_prefix(self, monkeypatch): + """load_weights raises on unknown non-codec checkpoint prefix.""" + talker, weights, mod, fake_load = self._make_fake_talker_and_keys() + monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load) + weights["unknown.prefix.weight"] = torch.randn(4) + with pytest.raises(ValueError, match="Unexpected checkpoint key"): + talker.load_weights(iter(weights.items())) + + +# ---- AC-7: Stage Input Processor ---- + + +class TestStageInputProcessor: + def test_revert_delay_pattern(self): + from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( + _revert_delay_pattern, + ) + + # 8 codebooks, 3 real frames -> delayed shape [8, 3+8-1=10] + Q, T = 8, 3 + delayed = torch.full((Q, T + Q - 1), 1024) # Fill with BOC + for i in range(Q): + for t in range(T): + delayed[i, i + t] = i * 100 + t # Real codes at shifted positions + result = _revert_delay_pattern(delayed) + assert result.shape == (Q, T) + for i in range(Q): + for t in range(T): + assert result[i, t].item() == i * 100 + t + + def test_revert_delay_pattern_rejects_wrong_codebooks(self): + from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( + _revert_delay_pattern, + ) + + # 7 codebooks should be rejected + codes = torch.zeros(7, 20) + with pytest.raises(ValueError, match="Expected exactly 8 codebook rows"): + _revert_delay_pattern(codes) + + def test_revert_delay_pattern_rejects_too_few_frames(self): + from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( + _revert_delay_pattern, + ) + + # 8 codebooks but only 5 frames (need at least 8) + codes = torch.zeros(8, 5) + with pytest.raises(ValueError, match="Not enough frames"): + _revert_delay_pattern(codes) + + def test_talker2code2wav_skips_too_few_frames_without_crashing(self): + from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( + talker2code2wav, + ) + + class Output: + multimodal_output = {"codes": {"audio": torch.zeros(5, 8, dtype=torch.long)}} + + class TalkerOutput: + finished = True + outputs = [Output()] + + result = talker2code2wav([TalkerOutput()]) + + assert len(result) == 1 + assert result[0]["prompt_token_ids"] == [] + + def test_filter_real_code_frames(self): + from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( + _filter_real_code_frames, + ) + + # 8 codebooks, 4 frames + codes = torch.tensor( + [ + [100, 200, 1024, 300], # cb0: frame 2 has BOC + [101, 201, 1024, 301], + [102, 202, 1024, 302], + [103, 203, 1024, 303], + [104, 204, 1024, 304], + [105, 205, 1024, 305], + [106, 206, 1024, 306], + [107, 207, 1024, 307], + ] + ) + result = _filter_real_code_frames(codes) + # Frame 2 (column 2) has BOC in all codebooks -> filtered out + assert result.shape == (8, 3) + assert result[0, 0].item() == 100 + assert result[0, 1].item() == 200 + assert result[0, 2].item() == 300 + + +# ---- AC-10: Registry ---- + + +class TestRegistry: + def test_talker_registered(self): + from vllm_omni.model_executor.models.registry import _OMNI_MODELS + + assert "HiggsMultimodalQwen3ForConditionalGeneration" in _OMNI_MODELS + assert "HiggsAudioV3TalkerForConditionalGeneration" in _OMNI_MODELS + + def test_code2wav_registered(self): + from vllm_omni.model_executor.models.registry import _OMNI_MODELS + + assert "HiggsAudioV3Code2WavForConditionalGeneration" in _OMNI_MODELS + + def test_pipeline_registered(self): + from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES + + assert "higgs_multimodal_qwen3" in _OMNI_PIPELINES + + def test_deploy_yaml_exists(self): + import os + + # __file__ = tests/unit/higgs_audio_v3/test_*.py → 4 dirnames to repo root + repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + deploy_dir = os.path.join(repo_root, "vllm_omni", "deploy") + for name in ( + "higgs_multimodal_qwen3.yaml", + "higgs_multimodal_qwen3_high_throughput.yaml", + "higgs_multimodal_qwen3_low_latency.yaml", + ): + yaml_path = os.path.join(deploy_dir, name) + assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}" + + +# ---- AC-3: Prompt Builder ---- + + +class TestPromptBuilder: + def _make_mock_tokenizer(self): + """Create a mock tokenizer with the required special tokens.""" + + class MockTokenizer: + def __init__(self): + self._added_vocab = { + "<|tts|>": 151700, + "<|text|>": 151701, + "<|audio|>": 151702, + "<|ref_audio|>": 151703, + "<|ref_text|>": 151704, + } + + def get_added_vocab(self): + return self._added_vocab + + def encode(self, text, add_special_tokens=True): + # Simple word-level tokenization for testing + return list(range(100, 100 + len(text.split()))) + + return MockTokenizer() + + def test_plain_tts_prompt(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + HiggsAudioV3TokenizerAdapter, + ) + + tok = self._make_mock_tokenizer() + adapter = HiggsAudioV3TokenizerAdapter(tok) + ids = adapter.build_prompt("Hello world") + assert ids[0] == 151700 # <|tts|> + assert ids[1] == 151701 # <|text|> + assert ids[-1] == 151702 # <|audio|> + assert len(ids) == 2 + 2 + 1 # tts + text + 2 word tokens + audio + + def test_empty_text_rejected(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + HiggsAudioV3TokenizerAdapter, + ) + + tok = self._make_mock_tokenizer() + adapter = HiggsAudioV3TokenizerAdapter(tok) + with pytest.raises(ValueError, match="non-empty"): + adapter.build_prompt("") + + def test_whitespace_only_rejected(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + HiggsAudioV3TokenizerAdapter, + ) + + tok = self._make_mock_tokenizer() + adapter = HiggsAudioV3TokenizerAdapter(tok) + with pytest.raises(ValueError, match="non-empty"): + adapter.build_prompt(" ") + + def test_missing_specials_rejected(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + HiggsAudioV3TokenizerAdapter, + ) + + class BadTokenizer: + def get_added_vocab(self): + return {"<|tts|>": 1} # Missing <|text|> and <|audio|> + + with pytest.raises(ValueError, match="missing"): + HiggsAudioV3TokenizerAdapter(BadTokenizer()) + + def test_no_voice_clone_tokens_in_plain_tts(self): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + HiggsAudioV3TokenizerAdapter, + ) + + tok = self._make_mock_tokenizer() + adapter = HiggsAudioV3TokenizerAdapter(tok) + ids = adapter.build_prompt("Hello") + # Should not contain ref_audio or ref_text token IDs + assert 151703 not in ids # <|ref_audio|> + assert 151704 not in ids # <|ref_text|> + + +class TestVoiceCloneReferenceCache: + def test_audio_tokenizer_dir_env_accepts_parent_or_subdir(self, tmp_path, monkeypatch): + from vllm_omni.model_executor.models.higgs_audio_v2 import higgs_audio_v2_tokenizer as tok + + parent = tmp_path / "OmniVoice" + subdir = parent / "audio_tokenizer" + subdir.mkdir(parents=True) + (subdir / "config.json").write_text("{}", encoding="utf-8") + + monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent)) + assert tok._resolve_audio_tokenizer_dir() == str(subdir) + + monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir)) + assert tok._resolve_audio_tokenizer_dir() == str(subdir) + + def test_higgs_v3_ref_code_cache_returns_clone(self): + from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + + serving = object.__new__(OmniOpenAIServingSpeech) + serving._higgs_audio_v3_ref_code_cache = OrderedDict() + serving._higgs_audio_v3_ref_code_cache_bytes = 0 + serving._higgs_audio_v3_ref_code_inflight = {} + + codes = torch.arange(16, dtype=torch.long).reshape(2, 8) + serving._put_higgs_audio_v3_ref_codes("ref-a", codes) + + cached = serving._get_higgs_audio_v3_ref_codes("ref-a") + assert cached is not None + cached.fill_(0) + + cached_again = serving._get_higgs_audio_v3_ref_codes("ref-a") + assert cached_again is not None + assert torch.equal(cached_again, codes) + + def test_higgs_v3_ref_code_inflight_deduplicates_concurrent_encode(self): + from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + + serving = object.__new__(OmniOpenAIServingSpeech) + serving._higgs_audio_v3_ref_code_cache = OrderedDict() + serving._higgs_audio_v3_ref_code_cache_bytes = 0 + serving._higgs_audio_v3_ref_code_inflight = {} + calls = 0 + + def encode_reference_audio(_wav, _sr): + nonlocal calls + calls += 1 + time.sleep(0.05) + return torch.arange(16, dtype=torch.long).reshape(2, 8) + + def apply_delay_pattern(codes): + return codes + 1 + + async def run(): + return await asyncio.gather( + *[ + serving._resolve_higgs_audio_v3_ref_codes( + "ref-a", + object(), + 24000, + encode_reference_audio, + apply_delay_pattern, + ) + for _ in range(3) + ] + ) + + results = asyncio.run(run()) + assert calls == 1 + assert sum(int(inflight_wait) for _, _, inflight_wait in results) == 2 + for codes, cache_hit, _ in results: + assert cache_hit is False + assert torch.equal(codes, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1) + + cached, cache_hit, inflight_wait = asyncio.run( + serving._resolve_higgs_audio_v3_ref_codes( + "ref-a", + object(), + 24000, + encode_reference_audio, + apply_delay_pattern, + ) + ) + assert calls == 1 + assert cache_hit is True + assert inflight_wait is False + assert torch.equal(cached, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1) diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index ae6ecc57c..6f4994b2c 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -22,6 +22,7 @@ class DummyInputBatch: def __init__(self, req_ids): self.req_ids = req_ids + self.req_id_to_index = {r: i for i, r in enumerate(req_ids)} class DummyReqState: diff --git a/vllm_omni/assets/video.py b/vllm_omni/assets/video.py deleted file mode 100644 index 6a5f3204a..000000000 --- a/vllm_omni/assets/video.py +++ /dev/null @@ -1,16 +0,0 @@ -import numpy as np -from vllm.assets.video import VideoAsset -from vllm.multimodal.media.audio import load_audio - - -def extract_video_audio(path: str = None, sampling_rate: int = 16000) -> np.ndarray: - """This function extracts the audio from a video file path and returns the audio as a numpy array. - Args: - path: The path to the video file. - Returns: - The audio as a numpy array. - """ - if not path: - path = VideoAsset(name="baby_reading").video_path - audio_signal, sr = load_audio(path, sr=sampling_rate) - return audio_signal diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 3357fa8bc..76b8628c1 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -16,8 +16,8 @@ from datetime import datetime from typing import Any, Literal import aiohttp +import numpy as np import pybase64 as base64 -from pydub import AudioSegment from tqdm.asyncio import tqdm from vllm.benchmarks import datasets from vllm.benchmarks.datasets import SampleRequest @@ -653,9 +653,8 @@ async def async_request_openai_chat_omni_completions( # Reset per-attempt state so that retries do not mix partial # outputs or metrics from previous attempts. generated_text = "" - generated_audio = None # For wav responses, accumulate decoded PCM bytes per chunk - # to avoid repeated AudioSegment decode/concat. + # to avoid repeated decode/concat. wav_pcm_buffer = bytearray() wav_audio_params: tuple[int, int, int] | None = None wav_inconsistent_chunk_count = 0 @@ -795,30 +794,33 @@ async def async_request_openai_chat_omni_completions( output.latency = timestamp - st output.generated_text = generated_text + audio_duration_sec = 0.0 + audio_frames = 0 if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None: channels, sample_width, frame_rate = wav_audio_params - generated_audio = AudioSegment( - data=bytes(wav_pcm_buffer), - sample_width=sample_width, - frame_rate=frame_rate, - channels=channels, - ) + frame_width = sample_width * channels + if frame_width > 0: + audio_frames = len(wav_pcm_buffer) // frame_width + audio_duration_sec = audio_frames / frame_rate + else: + logger.warning("Audio frame width is zero") elif audio_bytes_buffer: try: - generated_audio = AudioSegment.from_file( + from vllm.multimodal.audio import get_audio_duration + from vllm.multimodal.media.audio import load_audio + + waveform, sr = load_audio( io.BytesIO(bytes(audio_bytes_buffer)), - format=response_format, + sr=None, + mono=False, ) + audio_duration_sec = get_audio_duration(y=waveform, sr=sr) + audio_frames = int(audio_duration_sec * sr) except Exception as ex: logger.warning("Failed to decode accumulated audio bytes: %s", ex) - if generated_audio is not None: - output.audio_duration = len(generated_audio) / 1000.0 - frame_width = generated_audio.frame_width - if frame_width > 0: - output.audio_frames = len(generated_audio.raw_data) // frame_width - else: - output.audio_frames = 0 - logger.warning("Audio frame width is zero") + if audio_duration_sec > 0 or audio_frames > 0: + output.audio_duration = audio_duration_sec + output.audio_frames = audio_frames audio_duration = output.audio_duration if audio_duration > 0: output.audio_rtf = audio_generate_time / output.audio_duration @@ -827,8 +829,29 @@ async def async_request_openai_chat_omni_completions( logger.warning("Audio duration is zero") if _seed_tts_capture_pcm_for_wer() and getattr(request_func_input, "seed_tts_row", False): try: - seg = generated_audio.set_frame_rate(24000).set_channels(1).set_sample_width(2) - output.tts_output_pcm_bytes = bytes(seg.raw_data) + if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None: + from vllm.multimodal.audio import AudioResampler + + pcm_channels, pcm_sw, pcm_rate = wav_audio_params + pcm = np.frombuffer( + bytes(wav_pcm_buffer), dtype=np.int16 if pcm_sw == 2 else np.float32 + ) + if pcm_channels > 1: + pcm = pcm.reshape(-1, pcm_channels).mean(axis=1).astype(pcm.dtype) + pcm_f32 = pcm.astype(np.float32) / 32767.0 if pcm.dtype == np.int16 else pcm + if pcm_rate != 24000: + resampler = AudioResampler(target_sr=24000) + pcm_f32 = resampler.resample(pcm_f32, orig_sr=pcm_rate) + output.tts_output_pcm_bytes = (pcm_f32 * 32767).astype(np.int16).tobytes() + elif audio_bytes_buffer: + from vllm.multimodal.media.audio import load_audio + + waveform, _ = load_audio( + io.BytesIO(bytes(audio_bytes_buffer)), + sr=24000, + mono=True, + ) + output.tts_output_pcm_bytes = (waveform * 32767).astype(np.int16).tobytes() except Exception as ex: logger.warning("seed_tts WER PCM export failed: %s", ex) output.success = True diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index f02c07588..e5762456f 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -6,7 +6,6 @@ from vllm_omni.config.lora import LoRAConfig from vllm_omni.config.model import OmniModelConfig from vllm_omni.config.stage_config import ( DeployConfig, - ModelPipeline, PipelineConfig, StageConfig, StageConfigFactory, @@ -30,7 +29,6 @@ __all__ = [ "LoRAConfig", "StageConfig", "StageConfigFactory", - "ModelPipeline", "StageType", "StageExecutionType", "StagePipelineConfig", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 28ff5267a..106b2bf12 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -101,6 +101,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.mimo_audio.pipeline", "MIMO_AUDIO_PIPELINE", ), + "ming_tts": ( + "vllm_omni.model_executor.models.ming_tts.pipeline", + "MING_TTS_PIPELINE", + ), "voxtral_tts": ( "vllm_omni.model_executor.models.voxtral_tts.pipeline", "VOXTRAL_TTS_PIPELINE", @@ -149,6 +153,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.higgs_audio_v2.pipeline", "HIGGS_AUDIO_V2_PIPELINE", ), + "higgs_multimodal_qwen3": ( + "vllm_omni.model_executor.models.higgs_audio_v3.pipeline", + "HIGGS_AUDIO_V3_PIPELINE", + ), "dynin_omni": ( "vllm_omni.model_executor.models.dynin_omni.pipeline", "DYNIN_OMNI_PIPELINE", diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index ed1d1ca18..1bff5ef94 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -20,13 +20,6 @@ from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler -_MODELS_DIR = Path(__file__).resolve().parent.parent / "model_executor" / "models" - - -def get_pipeline_path(model_dir: str, filename: str) -> Path: - return _MODELS_DIR / model_dir / filename - - logger = init_logger(__name__) @@ -1061,77 +1054,6 @@ class StageConfig: return create_config(config_dict) -@dataclass -class ModelPipeline: - """Complete pipeline definition for a multi-stage model (legacy). - - TODO(@lishunyang12): remove once all models migrate to PipelineConfig. - """ - - model_type: str - stages: list[StageConfig] - - # Pipeline-wide behavior flags - async_chunk: bool = False - - # Optional distributed configuration - connectors: dict[str, Any] | None = None - edges: list[dict[str, Any]] | None = None - - def get_stage(self, stage_id: int) -> StageConfig | None: - """Look up a stage by its ID. - - Args: - stage_id: The stage ID to search for. - - Returns: - The matching StageConfig, or None if not found. - """ - for stage in self.stages: - if stage.stage_id == stage_id: - return stage - return None - - def validate_pipeline(self) -> list[str]: - """Validate pipeline topology at model integration time (not runtime). - - Checks: - - All stage IDs are unique - - All input_sources reference valid stage IDs - - At least one entry point (stage with empty input_sources) - - Returns: - List of validation error messages. Empty list if valid. - """ - errors: list[str] = [] - - if not self.stages: - errors.append("Topology has no stages defined") - return errors - - # Check for unique stage IDs - stage_ids = [s.stage_id for s in self.stages] - if len(stage_ids) != len(set(stage_ids)): - errors.append("Duplicate stage IDs found") - - stage_id_set = set(stage_ids) - - # Check input_sources reference valid stages - for stage in self.stages: - for source_id in stage.input_sources: - if source_id not in stage_id_set: - errors.append(f"Stage {stage.stage_id} references non-existent input source {source_id}") - if source_id == stage.stage_id: - errors.append(f"Stage {stage.stage_id} references itself as input source") - - # Check for at least one entry point - entry_points = [s for s in self.stages if not s.input_sources] - if not entry_points: - errors.append("No entry point found (stage with empty input_sources)") - - return errors - - class StageConfigFactory: """Factory that loads pipeline YAML and merges CLI overrides. @@ -1206,33 +1128,22 @@ class StageConfigFactory: continue return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path) - # --- Legacy path: load from pipeline YAML --- - pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code) + # An explicit deploy config can select a pipeline when the model's HF + # metadata is too generic for registry auto-detection. + if deploy_config_path is not None: + deploy_path = Path(deploy_config_path) + if deploy_path.exists(): + deploy_config = load_deploy_config(deploy_path) + if deploy_config.pipeline in _PIPELINE_REGISTRY: + return cls._create_from_registry( + deploy_config.pipeline, + cli_overrides, + deploy_config_path, + ) - if pipeline is None: - return None - - errors = pipeline.validate_pipeline() - if errors: - logger.warning(f"Pipeline validation warnings for {model}: {errors}") - - # Materialize the resolved pipeline-wide async_chunk value into every - # stage so build_engine_args_dict() can inject the stage connector - # spec and explicit False overrides are preserved. - resolved_async_chunk = cli_overrides.get("async_chunk") - if resolved_async_chunk is None: - resolved_async_chunk = bool(pipeline.async_chunk) - for stage in pipeline.stages: - stage.yaml_engine_args["async_chunk"] = bool(resolved_async_chunk) - - # Apply CLI overrides - result: list[StageConfig] = [] - for stage in pipeline.stages: - # Merge global CLI overrides - stage.runtime_overrides = cls._merge_cli_overrides(stage, cli_overrides) - result.append(stage) - - return result + # Not in the pipeline registry — let the caller fall back to the + # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path). + return None @classmethod def _create_from_registry( @@ -1355,193 +1266,6 @@ class StageConfigFactory: return [config_dict] - @classmethod - def _load_pipeline(cls, model: str, trust_remote_code: bool = True) -> ModelPipeline | None: - """Load a legacy ``pipeline.yaml`` for the model. - - Searches ``model_executor/models//pipeline.yaml`` by trying - (a) the raw ``model_type`` as the directory name, then - (b) ``model_type`` with hyphens replaced by underscores, - and finally (c) scanning every ``pipeline.yaml`` for one that - declares a matching ``model_type`` or ``hf_architectures``. - - Returns None if no pipeline.yaml is found — caller handles the - ``resolve_model_config_path`` fallback via stage_configs/ YAMLs. - """ - model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code) - if model_type is None: - return None - - # Direct lookups by convention - candidates = [model_type, model_type.replace("-", "_")] - for dir_name in candidates: - pipeline_path = get_pipeline_path(dir_name, "pipeline.yaml") - if pipeline_path.exists(): - return cls._parse_pipeline_yaml(pipeline_path, model_type) - - # Scan fallback: read every pipeline.yaml and match on declared fields - hf_archs = set(getattr(hf_config, "architectures", []) or []) if hf_config else set() - if _MODELS_DIR.exists(): - for subdir in sorted(_MODELS_DIR.iterdir()): - if not subdir.is_dir(): - continue - pipeline_path = subdir / "pipeline.yaml" - if not pipeline_path.exists(): - continue - try: - cfg = load_yaml_config(pipeline_path) - except Exception as exc: - logger.debug("Skip %s: %s", pipeline_path, exc) - continue - declared_type = getattr(cfg, "model_type", None) - declared_archs = set(getattr(cfg, "hf_architectures", None) or []) - if declared_type == model_type or (hf_archs and hf_archs.intersection(declared_archs)): - return cls._parse_pipeline_yaml(pipeline_path, declared_type or model_type) - - logger.debug("No pipeline.yaml found for model_type %s (archs=%s)", model_type, sorted(hf_archs)) - return None - - # Keys consumed as explicit StageConfig fields — everything else is - # passed through via yaml_extras. - _KNOWN_STAGE_KEYS: set[str] = { - "stage_id", - "model_stage", - "stage_type", - "input_sources", - "engine_input_source", - "custom_process_input_func", - "final_output", - "final_output_type", - "worker_type", - "scheduler_cls", - "hf_config_name", - "is_comprehension", - "engine_args", - "runtime", - } - - @classmethod - def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline: - """Parse a pipeline YAML file. - - Args: - path: Path to the YAML file. - model_type: Model type identifier. - - Returns: - ModelPipeline object. - """ - config_data = load_yaml_config(path) - - stages: list[StageConfig] = [] - for stage_data in config_data.stages: - # Use .get() for optional fields — idiomatic for OmegaConf DictConfig - stage_type_str = stage_data.get("stage_type", "llm") - stage_type = StageType(stage_type_str) if stage_type_str else StageType.LLM - - # Handle both 'input_sources' (new) and 'engine_input_source' (legacy) - input_sources = stage_data.get("input_sources", None) - if input_sources is None: - input_sources = stage_data.get("engine_input_source", []) - if input_sources is None: - input_sources = [] - input_sources = list(input_sources) - - # Extract per-stage engine_args and runtime dicts - raw_ea = stage_data.get("engine_args", None) - yaml_engine_args = to_dict(raw_ea) if raw_ea is not None else {} - raw_rt = stage_data.get("runtime", None) - yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {} - - # Migrate legacy runtime.max_batch_size → engine_args.max_num_seqs - if "max_batch_size" in yaml_runtime: - mbs = yaml_runtime.pop("max_batch_size") - yaml_engine_args.setdefault("max_num_seqs", int(mbs)) - logger.debug( - "Stage %s: migrated runtime.max_batch_size=%s to engine_args.max_num_seqs", - stage_data.get("stage_id", "?"), - mbs, - ) - - # Topology-level fields that also live inside engine_args in legacy - # YAMLs (worker_type, scheduler_cls, etc.) — read from both places. - worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None) - scheduler_cls = stage_data.get("scheduler_cls", None) or yaml_engine_args.pop("scheduler_cls", None) - if scheduler_cls: - async_sched = yaml_engine_args.get("async_scheduling") - if async_sched is not None: - logger.warning( - "Stage %s: async_scheduling=%r and scheduler_cls=%r " - "should not be set together. scheduler_cls will take " - "precedence for which scheduler is used.", - stage_data.stage_id, - async_sched, - scheduler_cls, - ) - else: - logger.warning( - "Stage %s: scheduler_cls=%r is deprecated. Use async_scheduling instead.", - stage_data.stage_id, - scheduler_cls, - ) - hf_config_name = stage_data.get("hf_config_name", None) or yaml_engine_args.pop("hf_config_name", None) - model_stage = getattr(stage_data, "model_stage", None) or yaml_engine_args.pop("model_stage", None) - - # Collect pass-through fields (default_sampling_params, - # output_connectors, input_connectors, tts_args, etc.) - yaml_extras: dict[str, Any] = {} - for key in stage_data: - if key not in cls._KNOWN_STAGE_KEYS: - val = stage_data[key] - try: - yaml_extras[key] = to_dict(val) - except ValueError: - yaml_extras[key] = val - - stage = StageConfig( - stage_id=stage_data.stage_id, - model_stage=model_stage or "", - stage_type=stage_type, - input_sources=input_sources, - custom_process_input_func=stage_data.get("custom_process_input_func", None), - final_output=stage_data.get("final_output", False), - final_output_type=stage_data.get("final_output_type", None), - worker_type=worker_type, - scheduler_cls=scheduler_cls, - hf_config_name=hf_config_name, - is_comprehension=stage_data.get("is_comprehension", False), - yaml_engine_args=yaml_engine_args, - yaml_runtime=yaml_runtime, - yaml_extras=yaml_extras, - ) - stages.append(stage) - - # Get pipeline-wide flags - async_chunk = config_data.get("async_chunk", False) - - # Get optional connector config — check both top-level and nested - # under ``runtime`` (legacy stage_configs format). - connectors = None - edges = None - if hasattr(config_data, "connectors"): - connectors = to_dict(config_data.connectors) - if hasattr(config_data, "edges"): - edges = to_dict(config_data.edges) - if hasattr(config_data, "runtime") and config_data.runtime is not None: - top_runtime = config_data.runtime - if connectors is None and hasattr(top_runtime, "connectors"): - connectors = to_dict(top_runtime.connectors) - if edges is None and hasattr(top_runtime, "edges"): - edges = to_dict(top_runtime.edges) - - return ModelPipeline( - model_type=getattr(config_data, "model_type", model_type), - stages=stages, - async_chunk=async_chunk, - connectors=connectors, - edges=edges, - ) - @classmethod def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]: """Auto-detect model_type from model directory. diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py index 79f0990fd..25acfe8ed 100644 --- a/vllm_omni/core/prefix_cache.py +++ b/vllm_omni/core/prefix_cache.py @@ -391,27 +391,33 @@ class OmniTensorPrefixCache: self, query_start_loc: torch.Tensor, input_batch: InputBatch, - multimodal_outputs: dict, + multimodal_outputs: dict | None, num_scheduled_tokens: dict[str, int], ): """Get the merged multimodal states if hidden state prefix caching is enabled.""" combined_multimodal_outputs = {} + # Talkers that produce multimodal outputs only at decode (e.g. Higgs + # Audio v3's audio codes) have ``multimodal_outputs is None`` at the + # initial prefill step. Treat it as an empty mapping so the merger + # short-circuits cleanly and the cache merge still runs once codes + # start arriving. + mm_outputs = multimodal_outputs if multimodal_outputs is not None else {} # First get the prefix cached tensors that are present in the mm data for mm_key in self.mm_cache_keys: - if mm_key in multimodal_outputs: + if mm_key in mm_outputs: combined_multimodal_outputs[mm_key] = self._get_merged_tensors( query_start_loc=query_start_loc, input_batch=input_batch, cache=self.mm_outputs_cache[mm_key], - hidden_states=multimodal_outputs[mm_key], + hidden_states=mm_outputs[mm_key], num_scheduled_tokens=num_scheduled_tokens, ) # Then, get everything else (passthrough data); first, convert to CPU # tensors similarly to the non prefix cached path, and then populate # the subdicts mapping request IDs -> payload objects - passthrough_keys = set(multimodal_outputs.keys()) - self.mm_cache_keys - passthrough_mm_data = {k: v for k, v in multimodal_outputs.items() if k in passthrough_keys} + passthrough_keys = set(mm_outputs.keys()) - self.mm_cache_keys + passthrough_mm_data = {k: v for k, v in mm_outputs.items() if k in passthrough_keys} mm_cpu = build_mm_cpu(multimodal_outputs=passthrough_mm_data) for mm_key, mm_val in mm_cpu.items(): diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index a6d5edd8f..4c4ae361a 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -13,7 +13,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMSchedu from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.core.sched.utils import remove_all -from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs +from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs, FinishReason from vllm.v1.metrics.perf import PerfStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate @@ -69,6 +69,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # alive. Stopped on the first scheduler step after extraction ack. self.pending_stop_after_extraction: set[str] = set() + self.finished_req_ids_dict = defaultdict(set) + # [Omni] Pre-parse KV transfer criteria self.kv_transfer_criteria = self._get_kv_transfer_criteria() @@ -216,7 +218,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self._process_pending_input_timeouts() if self.chunk_transfer_adapter: - self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) + self.chunk_transfer_adapter.process_pending_chunks( + self.waiting, self.running, scheduler_requests=self.requests + ) try: scheduler_output = super().schedule() @@ -556,16 +560,24 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # outputs in this step. engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()} + # FIXME: finished_req_ids_dict is unconditionally initialized as + # defaultdict(set) in __init__ (not gated by include_finished_set). + # This branch is therefore always eligible once any client_index is + # populated; revisit when wiring streaming-only / upstream semantics. finished_req_ids = self.finished_req_ids_dict if finished_req_ids: # Include ids of requests that finished since last outputs # were sent. for client_index, finished_set in finished_req_ids.items(): - # Set finished request set in EngineCoreOutputs for this client. - if (eco := engine_core_outputs.get(client_index)) is not None: - eco.finished_requests = finished_set - else: - engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set) + eco = engine_core_outputs.get(client_index) + if eco is None: + eco = EngineCoreOutputs() + engine_core_outputs[client_index] = eco + emitted = {o.request_id for o in eco.outputs} + for req_id in finished_set: + if req_id not in emitted: + eco.outputs.append(EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT)) + eco.finished_requests = finished_set finished_req_ids.clear() if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None: @@ -616,7 +628,34 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + # Realign stale ``request.status`` (chunk-transfer-adapter's + # ``requests_origin_status`` table doesn't follow the + # ``waiting → running`` admit transition; without this, an abort + # arriving between admit and the next deque round-trip leaves + # the request in ``self.running`` with ``status=WAITING`` and + # upstream ``Scheduler.finish_requests`` silently fails to + # release the worker's ``input_batch`` slot -- after + # ``max_num_seqs`` such aborts new requests hang at + # ``chunks=0``). Only the ``async_chunk`` path triggers the + # staleness; with ``async_chunk`` disabled this is a cheap O(n) + # no-op over an already-aligned set, kept unconditional so the + # abort path stays uniform across configurations. See + # ``OmniSchedulerMixin._realign_request_status_to_queues`` and + # #3774 discussion. + self._realign_request_status_to_queues(request_ids) + finished = super().finish_requests(request_ids, finished_status) + + # Defensive post-finish purge: belt-and-suspenders to the + # realignment above. Even after realign + ``super()``, corner + # cases (mid-transition status, connector cleanups that pop + # from ``self.requests`` without unwinding ``self.running``) + # can leave already-finished or untracked entries in + # ``self.running``. Sweep them now so the worker's + # ``input_batch`` slot never pins a freed request and starves + # new admissions. See ``OmniSchedulerMixin._purge_finished_from_running``. + self._purge_finished_from_running() + input_coordinator = getattr(self, "input_coordinator", None) if input_coordinator is not None: for request_id, _ in finished: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 1c43b958a..9f5d3e104 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -87,7 +87,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._consume_pending_connector_output(model_mode="generation") self._process_pending_input_timeouts() if self.chunk_transfer_adapter: - self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) + self.chunk_transfer_adapter.process_pending_chunks( + self.waiting, self.running, scheduler_requests=self.requests + ) # OMNI: Track requests that are already finished (e.g., marked by connector) # These should be removed from running and not scheduled @@ -364,7 +366,23 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + # See ``OmniSchedulerMixin._realign_request_status_to_queues`` -- + # closes the residual hang on the + # ``waiting → running → abort-before-next-deque-round-trip`` race + # surfaced by #3774's reproduction matrix. Only the + # ``async_chunk`` path actually triggers the staleness; with + # ``async_chunk`` disabled this is a cheap O(n) no-op, kept + # unconditional so the abort path stays uniform. + self._realign_request_status_to_queues(request_ids) + finished = super().finish_requests(request_ids, finished_status) + + # See ``OmniSchedulerMixin._purge_finished_from_running`` -- + # defensive belt-and-suspenders sweep paired with the realign + # above. Closes residual ``self.running`` slot leaks if any + # corner case slips past upstream's status-driven removal. + self._purge_finished_from_running() + if self.input_coordinator is not None: for request_id, _ in finished: self._free_input_coordinator_request(request_id) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index c33d16dac..165c04e96 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -2,6 +2,7 @@ from __future__ import annotations import os import time +from collections.abc import Iterable from typing import Any from vllm.logger import init_logger @@ -156,3 +157,122 @@ class OmniSchedulerMixin: return None self._last_stats_time = now return super().make_stats(*args, **kwargs) + + def _realign_request_status_to_queues( + self, + request_ids: str | Iterable[str] | None, + ) -> None: + """Realign ``request.status`` to actual queue membership. + + ``OmniChunkTransferAdapter._process_chunk_queue`` stamps + ``requests_origin_status[req.id] = WAITING`` (or ``RUNNING``) when + first parking a request in a chunk-transfer deque. On the next + tick, when the chunk arrives, ``_process_chunk_queue`` sets + ``request.status = target_status`` and continues, but + ``requests_origin_status`` is left at its first-park value -- no + hook updates it on the ``waiting → running`` admit transition + that ``super().schedule()`` later performs. The table stays + stale until the request makes another deque round-trip. + + If an abort lands in the gap between admit and the next deque + round-trip, ``chunk_transfer_adapter.finish_requests`` reads the + stale ``WAITING`` from ``requests_origin_status``, stomps it + onto ``request.status``, and the upstream + ``Scheduler.finish_requests`` else branch silently fails to + remove from ``self.running`` -- the request stays alive in + ``self.running`` and the worker's ``input_batch`` slot leaks. + After ``max_num_seqs`` such aborts every new request hangs at + ``chunks=0`` until the client times out. + + Realign here: if a request lives in ``self.running`` but its + status is not ``RUNNING``, set it to ``RUNNING``; symmetrically + flip ``RUNNING → WAITING`` when the request is actually in + ``self.waiting``. This is a localized safety net for + ``requests_origin_status`` staleness on the admit transition; + it does not touch the adapter's invariants and is complementary + to the chunk-transfer-adapter deque purge that already runs + inside ``process_pending_chunks`` / ``restore_queues``. + + Note on scope: only the ``async_chunk`` path actually triggers + the ``requests_origin_status`` staleness this helper repairs. + When ``async_chunk`` is disabled, no chunk-transfer round-trip + occurs between admit and finish, so the realignment walk is a + cheap O(n) no-op over an already-aligned set. The call is kept + unconditional in ``finish_requests`` to (a) keep the abort path + uniform and (b) defend any future configuration that re-enables + chunk transfer from rediscovering the same regression. + + See https://github.com/vllm-project/vllm-omni/pull/3774 and the + residual-hang reproduction discussed in that PR. + """ + # Mirror the upstream Scheduler.finish_requests resolution of + # ``request_ids`` so realignment touches exactly the set that + # ``super().finish_requests`` will then walk. + if isinstance(request_ids, str): + ids_to_align: Iterable[str] = (request_ids,) + elif request_ids is None: + ids_to_align = list(self.requests.keys()) + else: + ids_to_align = list(request_ids) + + if not ids_to_align: + return + + running_ids = {r.request_id for r in self.running} + waiting_ids = {r.request_id for r in self.waiting} + + for rid in ids_to_align: + req = self.requests.get(rid) + if req is None or req.is_finished(): + continue + if rid in running_ids and req.status != RequestStatus.RUNNING: + req.status = RequestStatus.RUNNING + elif rid in waiting_ids and req.status == RequestStatus.RUNNING: + req.status = RequestStatus.WAITING + + def _purge_finished_from_running(self) -> None: + """Defensive post-finish sweep of ``self.running``. + + Belt-and-suspenders to ``_realign_request_status_to_queues``: + even after status realignment lets upstream + ``Scheduler.finish_requests`` pick the right removal branch, + a future regression or an unexpected ``status`` mid-transition + could still leave already-finished entries in ``self.running``. + Sweeping here guarantees the worker's ``input_batch`` slot is + not pinned by a freed request. + + Complementary to ``_realign_request_status_to_queues``: realign + is preventive (fix ``status`` before ``super().finish_requests`` + so the right branch fires); this purge is defensive (sweep the + residue after ``super().finish_requests`` so any stale entries + are reclaimed). + + Scope of the predicate. ``is_finished()`` covers entries the + upstream ``finish_requests`` already drained from ``self.requests`` + but failed to remove from ``self.running``; the + ``request_id not in self.requests`` arm catches the same surface + from a different angle and is the post-cleanup mirror of the + deque purge ``_purge_untracked_chunk_requests`` already runs at + the chunk-transfer-adapter layer. It does **not** by itself make + arbitrary direct deletions of ``self.requests`` safe -- callers + that pop ``self.requests`` outside the standard finish path + still have to go through ``_free_request`` (or equivalent) for + block / connector / coordinator cleanup. This sweep only + reclaims the ``self.running`` slot reference. + + In-place via ``self.running[:] = ...`` for minor consistency + with idiomatic vLLM scheduler mutation; upstream + ``Scheduler.finish_requests`` itself rebinds ``self.running``, + so list identity across the whole call is not preserved -- the + slice form is just to avoid an extra rebind inside this helper. + + Assumes the upstream V1 invariant that scheduler ticks are + serialized on a single thread; in-place mutation here is no more + racy than the rest of the scheduler under that assumption. + + See https://github.com/vllm-project/vllm-omni/pull/3774 + discussion. + """ + if not self.running: + return + self.running[:] = [req for req in self.running if not req.is_finished() and req.request_id in self.requests] diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index aa0e1005c..eec3e5c33 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -63,6 +63,7 @@ class OmniPayloadMeta(TypedDict, total=False): stream_finished: torch.Tensor req_id: list[str] left_context_size: int + right_holdback_size: int override_keys: list[tuple[str, str]] num_processed_tokens: int next_stage_prompt_len: int @@ -127,6 +128,7 @@ class EmbeddingsStruct(_StructBase): voice: torch.Tensor | None = None speech_feat: torch.Tensor | None = None speech_token: torch.Tensor | None = None + speech_token_len: torch.Tensor | None = None embedding: torch.Tensor | None = None thinker_reply: torch.Tensor | None = None @@ -150,6 +152,7 @@ class MetaStruct(_StructBase): stream_finished: torch.Tensor | None = None req_id: list[str] | None = None left_context_size: int | None = None + right_holdback_size: int | None = None override_keys: list[tuple[str, str]] | None = None num_processed_tokens: int | None = None next_stage_prompt_len: int | None = None diff --git a/vllm_omni/deploy/README_higgs_audio_v3.md b/vllm_omni/deploy/README_higgs_audio_v3.md new file mode 100644 index 000000000..ea1363069 --- /dev/null +++ b/vllm_omni/deploy/README_higgs_audio_v3.md @@ -0,0 +1,53 @@ +# Higgs-Audio V3 Deploy Profiles + +Higgs-Audio V3 has two Stage 0 graph profiles. They are intentionally separate +because the graph paths are mutually exclusive. + +## High Throughput + +Use `higgs_multimodal_qwen3_high_throughput.yaml` for medium/high concurrency +serving. This profile keeps Stage 0 `enforce_eager: true`, which preserves the +Higgs-specific local MLP CUDA graph path. It is the default production profile +for throughput-oriented serving. + +`higgs_multimodal_qwen3.yaml` is kept as the auto-discovered default deploy +config for `model_type=higgs_multimodal_qwen3`, and matches the high-throughput +profile. + +```bash +vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ + --omni --trust-remote-code \ + --deploy-config vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml +``` + +## Low Latency + +Use `higgs_multimodal_qwen3_low_latency.yaml` for low-concurrency serving +(for example c1-c4) where Stage 0 decode launch overhead dominates. This profile +sets Stage 0 `enforce_eager: false` and explicitly enables vLLM +`FULL_DECODE_ONLY` CUDA graph: + +```yaml +compilation_config: + cudagraph_capture_sizes: [1, 2, 4, 8, 16] + cudagraph_mode: FULL_DECODE_ONLY + cudagraph_num_of_warmups: 1 +``` + +FULL_DECODE is controlled by deploy configuration, not by an environment +variable. When this external decode graph is active, the Higgs talker disables +the local MLP CUDA graph automatically. + +```bash +vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ + --omni --trust-remote-code \ + --deploy-config vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml +``` + +## Notes + +- Stage 1 remains `enforce_eager: true` in both profiles. +- Keep `VLLM_USE_DEEP_GEMM=0` and `VLLM_MOE_USE_DEEP_GEMM=0` for this model + unless DeepGEMM support is revalidated. +- Revalidate end-to-end throughput and audio quality before changing the default + auto-discovered config. diff --git a/vllm_omni/deploy/cosyvoice3.yaml b/vllm_omni/deploy/cosyvoice3.yaml index 4bfd4ab85..d81cf4a9b 100644 --- a/vllm_omni/deploy/cosyvoice3.yaml +++ b/vllm_omni/deploy/cosyvoice3.yaml @@ -7,9 +7,22 @@ # references stay in the yaml unconditionally; the runtime only activates # them when ``async_chunk: true``. # -# enforce_eager=true everywhere: -# * Stage 0 (talker) — cudagraph not verified on this checkpoint. -# * Stage 1 (code2wav) — CUDA graphs don't work with dynamic conv shapes. +# enforce_eager: +# * Stage 0 (talker) — enforce_eager: false (CUDA graphs enabled). +# * Stage 1 (code2wav) — enforce_eager: true (CUDA graphs don't work with +# dynamic conv shapes). +# +# Streaming continuity is GPU-class sensitive — TUNE PER GPU. +# The defaults below (max_num_seqs: 8, codec_chunk_frames: 15) are tuned for +# H100, where streaming continuity at concurrency 4 stays ~100%. They push +# more audio per scheduler step through the SharedMemoryConnector than a +# slower GPU can keep up with: on H20-3e the same config regressed continuity +# to ~78% at c=4 (vs ~86% with smaller batch/chunk), because code2wav can't +# produce mel chunks fast enough to feed the streaming output → buffer +# underruns → audible gaps. +# On a slower GPU class, lower max_num_seqs (e.g. 4) and/or codec_chunk_frames +# first. Raising connector_get_sleep_s also relaxes the consumer's polling +# cadence, but trades against TTFP — keep it as a last resort. async_chunk: true dtype: float32 @@ -21,16 +34,16 @@ connectors: connector_get_sleep_s: 0.01 connector_get_max_wait_first_chunk: 3000 connector_get_max_wait: 300 - codec_chunk_frames: 25 + codec_chunk_frames: 15 codec_left_context_frames: 25 codec_vocab_size: 6561 stages: - stage_id: 0 max_num_batched_tokens: 32768 - max_num_seqs: 1 + max_num_seqs: 8 gpu_memory_utilization: 0.4 - enforce_eager: true + enforce_eager: false trust_remote_code: true enable_prefix_caching: false devices: "0" @@ -49,7 +62,7 @@ stages: - stage_id: 1 max_num_batched_tokens: 32768 - max_num_seqs: 1 + max_num_seqs: 8 gpu_memory_utilization: 0.2 enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml new file mode 100644 index 000000000..5fc44a295 --- /dev/null +++ b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml @@ -0,0 +1,81 @@ +# Higgs-audio v3 deploy: Stage 0 (talker) -> Stage 1 (code2wav) via shared memory. +# +# Auto-discovered default for model_type=higgs_multimodal_qwen3. This file is +# the high-throughput profile and intentionally keeps Stage 0 enforce_eager=true +# so the Higgs-specific local MLP CUDA graph path remains active. For low +# concurrency latency-sensitive serving, use +# higgs_multimodal_qwen3_low_latency.yaml explicitly. +# +# Plain TTS only (no voice cloning). async_chunk streams codec chunks from +# Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's +# worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size +# so the per-chunk codec windows stitch into a coherent PCM stream. +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + # Stream PCM bytes back to the client as Stage 1 emits each chunk. + codec_streaming: true + # Streaming chunk shape for talker2code2wav_async_chunk: + # codec_chunk_frames steady-state new frames per emit + # codec_left_context_frames sliding past context to avoid leading-edge artifacts + # codec_right_holdback_frames trailing frames held back so the codec + # has future context for emit-region tail samples + # initial_codec_chunk_frames smaller first chunk to minimize TTFA + codec_chunk_frames: 25 + codec_left_context_frames: 25 + codec_right_holdback_frames: 4 + initial_codec_chunk_frames: 1 + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 5000 + connector_get_max_wait: 500 + +stages: + - stage_id: 0 + max_num_seqs: 16 + gpu_memory_utilization: 0.6 + enforce_eager: true + attention_backend: FLASHINFER + trust_remote_code: true + # Stage-0 prefix caching is on by default. The talker's HS opt-out and + # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) + # keep per-step bookkeeping negligible; any shared text prefix across + # concurrent requests is a net throughput win. Stage 1 (codec decoder) + # has no KV cache and stays disabled. + enable_prefix_caching: true + async_scheduling: false + max_num_batched_tokens: 4096 + max_model_len: 8192 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 1.0 + top_p: 0.95 + top_k: 50 + max_tokens: 2048 + seed: 42 + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 16 + gpu_memory_utilization: 0.25 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml new file mode 100644 index 000000000..86427be13 --- /dev/null +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml @@ -0,0 +1,80 @@ +# Higgs-audio v3 deploy, high-throughput profile. +# +# This is the default production profile for medium/high concurrency serving. +# Stage 0 keeps enforce_eager=true so Higgs' local MLP CUDA graph remains active. +# Do not enable vLLM FULL_DECODE graph here: it replaces the local MLP graph +# path instead of stacking on top of it. +# +# Plain TTS only (no voice cloning). async_chunk streams codec chunks from +# Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's +# worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size +# so the per-chunk codec windows stitch into a coherent PCM stream. +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + # Stream PCM bytes back to the client as Stage 1 emits each chunk. + codec_streaming: true + # Streaming chunk shape for talker2code2wav_async_chunk: + # codec_chunk_frames steady-state new frames per emit + # codec_left_context_frames sliding past context to avoid leading-edge artifacts + # codec_right_holdback_frames trailing frames held back so the codec + # has future context for emit-region tail samples + # initial_codec_chunk_frames smaller first chunk to minimize TTFA + codec_chunk_frames: 25 + codec_left_context_frames: 25 + codec_right_holdback_frames: 4 + initial_codec_chunk_frames: 1 + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 5000 + connector_get_max_wait: 500 + +stages: + - stage_id: 0 + max_num_seqs: 64 + gpu_memory_utilization: 0.6 + enforce_eager: true + attention_backend: FLASHINFER + trust_remote_code: true + # Stage-0 prefix caching is on by default. The talker's HS opt-out and + # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) + # keep per-step bookkeeping negligible; any shared text prefix across + # concurrent requests is a net throughput win. Stage 1 (codec decoder) + # has no KV cache and stays disabled. + enable_prefix_caching: true + async_scheduling: false + max_num_batched_tokens: 4096 + max_model_len: 8192 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 1.0 + top_p: 0.95 + top_k: 50 + max_tokens: 2048 + seed: 42 + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 16 + gpu_memory_utilization: 0.25 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml new file mode 100644 index 000000000..8210db4a8 --- /dev/null +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml @@ -0,0 +1,85 @@ +# Higgs-audio v3 deploy, low-latency profile. +# +# Use this profile for low-concurrency serving (for example c1-c4) where Stage 0 +# decode launch overhead dominates. Stage 0 enables vLLM FULL_DECODE_ONLY CUDA +# graph through YAML configuration: +# +# enforce_eager: false +# compilation_config.cudagraph_mode: FULL_DECODE_ONLY +# +# This is not controlled by an environment variable. The Higgs talker disables +# its local MLP CUDA graph automatically when the external decode graph is active, +# because the two graph paths are mutually exclusive. +# +# Plain TTS only (no voice cloning). async_chunk streams codec chunks from +# Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's +# worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size +# so the per-chunk codec windows stitch into a coherent PCM stream. +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + # Stream PCM bytes back to the client as Stage 1 emits each chunk. + codec_streaming: true + # Streaming chunk shape for talker2code2wav_async_chunk: + # codec_chunk_frames steady-state new frames per emit + # codec_left_context_frames sliding past context to avoid leading-edge artifacts + # codec_right_holdback_frames trailing frames held back so the codec + # has future context for emit-region tail samples + # initial_codec_chunk_frames smaller first chunk to minimize TTFA + codec_chunk_frames: 25 + codec_left_context_frames: 25 + codec_right_holdback_frames: 4 + initial_codec_chunk_frames: 1 + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 5000 + connector_get_max_wait: 500 + +stages: + - stage_id: 0 + max_num_seqs: 16 + gpu_memory_utilization: 0.6 + enforce_eager: false + attention_backend: FLASHINFER + trust_remote_code: true + enable_prefix_caching: true + async_scheduling: false + compilation_config: + cudagraph_capture_sizes: [1, 2, 4, 8, 16] + cudagraph_mode: FULL_DECODE_ONLY + cudagraph_num_of_warmups: 1 + max_num_batched_tokens: 4096 + max_model_len: 8192 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 1.0 + top_p: 0.95 + top_k: 50 + max_tokens: 2048 + seed: 42 + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 16 + gpu_memory_utilization: 0.25 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 diff --git a/vllm_omni/deploy/hunyuan_image_3_moe.yaml b/vllm_omni/deploy/hunyuan_image_3_moe.yaml index a6e2add17..67ce8b66a 100644 --- a/vllm_omni/deploy/hunyuan_image_3_moe.yaml +++ b/vllm_omni/deploy/hunyuan_image_3_moe.yaml @@ -62,6 +62,7 @@ stages: max_tokens: 8192 detokenize: true skip_special_tokens: false + include_stop_str_in_output: true, - stage_id: 1 max_num_seqs: 1 diff --git a/vllm_omni/deploy/ming_tts.yaml b/vllm_omni/deploy/ming_tts.yaml new file mode 100644 index 000000000..1a2dc99ec --- /dev/null +++ b/vllm_omni/deploy/ming_tts.yaml @@ -0,0 +1,55 @@ +# Ming-omni-tts deploy: stage-0 LLM+flow -> stage-1 audio VAE. +# Verified legacy settings migrated from ming_tts.yaml and ming_tts_async_chunk.yaml. +# Default mode is async-chunk streaming on a single GPU. +# pipeline: required — HF config reports model_type="dense" with no architectures, +# so auto-detection falls through. This key routes the deploy-config probe in +# StageConfigFactory.create_from_model to the correct registry entry. +pipeline: ming_tts +async_chunk: true +trust_remote_code: false +dtype: bfloat16 +enable_prefix_caching: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + latent_chunk_size: 25 + latent_left_context: 0 + +stages: + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.45 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 512 + detokenize: true + + - stage_id: 1 + max_num_seqs: 1 + gpu_memory_utilization: 0.25 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + detokenize: false diff --git a/vllm_omni/deploy/moss_tts.yaml b/vllm_omni/deploy/moss_tts.yaml index acdb87721..dbab816eb 100644 --- a/vllm_omni/deploy/moss_tts.yaml +++ b/vllm_omni/deploy/moss_tts.yaml @@ -14,6 +14,7 @@ connectors: connector_get_max_wait: 300 codec_chunk_frames: 25 codec_left_context_frames: 0 + decode_cudagraph_capture_sizes: [4, 8, 16, 25, 32, 50, 64, 100, 128, 200, 256] stages: - stage_id: 0 @@ -38,7 +39,7 @@ stages: - stage_id: 1 max_num_seqs: 1 gpu_memory_utilization: 0.12 - enforce_eager: true + enforce_eager: false enable_prefix_caching: false async_scheduling: true # Talker can emit up to max_tokens=4096 frames × n_vq=32 = 131072 codec diff --git a/vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml b/vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml new file mode 100644 index 000000000..5c0d51e5a --- /dev/null +++ b/vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml @@ -0,0 +1,107 @@ +# Qwen3-Omni-MoE intra-node deploy using Mori XGMI transfer. +# +# Targets a single AMD Instinct MI300X OAM node (8x GPU, 192GB HBM3, +# gfx942). Each stage occupies one dedicated GPU (3 GPUs total) and +# stage-to-stage streaming chunks are shipped via Mori's RDMA IOEngine +# (with ``backend_type: xgmi`` selecting the AMD Infinity Fabric +# GPU-to-GPU direct path for intranode transfers). +# +# Pipeline wiring (stage_id, execution_type, model_stage, input_sources, +# custom processors, async-chunk builders) lives in +# ``vllm_omni/model_executor/models/qwen3_omni/pipeline.py`` and is +# merged into the stages below by the deploy loader. +# +# Why this yaml exists: +# * On CUDA/H100 Qwen3-Omni-MoE ships via ``SharedMemoryConnector`` +# (see ``qwen3_omni_moe.yaml``). On MI300X the SHM path still works +# but XGMI RDMA can substantially reduce stage-transition latency +# for the thinker->talker hidden-state stream and the talker->code2wav +# codec-frame stream, both of which run in async-chunk mode. +# * This config swaps the connector definition to +# ``MoriTransferEngineConnector`` while keeping the pipeline topology, +# sampling params, and async_chunk=true gating identical to the CUDA +# deploy, so the chunk_transfer_adapter path (scheduler-owned +# background put/get threads) is fully exercised. +# +# Usage: +# vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \ +# --deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml +# +# Dependencies: mori (``pip install mori``), torch-rocm>=2.9, vllm-omni +# with #2383 deploy schema, ibverbs/RDMA NICs exposed via /dev/infiniband. +# +# Limitation: TP=1-only. Connector raises NotImplementedError on +# tp_world_size > 1 (chunk-path endpoint derivation isn't rank-aware yet). +async_chunk: true + +connectors: + mori_connector: + name: MoriTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + backend_type: "xgmi" # AMD Infinity Fabric GPU-to-GPU direct + device_name: "" # leave empty to honour $MORI_RDMA_DEVICES + memory_pool_size: 536870912 # 512 MB, matches Qwen2.5-Omni Mori deploy + memory_pool_device: "cuda" # GPU memory required for XGMI RDMA transfers + xgmi_num_streams: 64 + xgmi_num_events: 64 + qp_per_transfer: 1 + num_worker_threads: 1 + post_batch_size: -1 + # Chunk-transfer knobs consumed by talker2code2wav_async_chunk: + # 25 decoded codec frames per connector.put (~40 ms at 625 Hz + # codec rate) with a 25-frame left context to satisfy the + # code2wav receptive field. Kept identical to the SHM deploy so + # accuracy stays comparable when swapping backends. + codec_chunk_frames: 25 + codec_left_context_frames: 25 + +stages: + - stage_id: 0 + gpu_memory_utilization: 0.9 + enforce_eager: true # ROCm override: avoid flashinfer autotune path + devices: "0" + output_connectors: + to_stage_1: mori_connector + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 42 + repetition_penalty: 1.05 + + - stage_id: 1 + gpu_memory_utilization: 0.9 + enforce_eager: true + devices: "1" + input_connectors: + from_stage_0: mori_connector + output_connectors: + to_stage_2: mori_connector + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.05 + + - stage_id: 2 + gpu_memory_utilization: 0.3 + max_num_seqs: 1 + enforce_eager: true + async_scheduling: false + # Codec prefill length (Q * num_frames) exceeds 32k default; + # matches qwen3_omni_moe.yaml code2wav sizing. + max_num_batched_tokens: 51200 + devices: "2" + input_connectors: + from_stage_1: mori_connector + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.1 diff --git a/vllm_omni/deploy/voxcpm2.yaml b/vllm_omni/deploy/voxcpm2.yaml index 4e3f7f3f6..5d366521b 100644 --- a/vllm_omni/deploy/voxcpm2.yaml +++ b/vllm_omni/deploy/voxcpm2.yaml @@ -1,5 +1,6 @@ # VoxCPM2 deploy: single-stage AR pipeline with per-request state batching. -# Verified on 1x H20 141GB. +# Runs on large (H20 141GB) and small (24GB L4) cards alike: the KV cache is +# pinned to a right-sized budget instead of ballooning to fill the card. # # VoxCPM2-local runtime knobs are passed through hf_overrides. async_chunk: false @@ -9,6 +10,16 @@ stages: - stage_id: 0 max_num_seqs: 32 gpu_memory_utilization: 0.9 + # Right-size the KV cache instead of letting it fill gpu_memory_utilization. + # max_num_seqs(32) * max_model_len(4096) ~= 131k tokens; 6 GiB holds ~175k + # (42x concurrency, above the 32-seq scheduler cap, so no throughput loss). + # The remaining VRAM stays free for the VoxCPM2 diffusion side-path + # (LocDiT / AudioVAE / code2wav + batch CUDA-graph pools), which allocates + # outside vLLM's KV accounting during inference. Keeps peak ~13 GiB on any + # card and avoids per-rebase OOM tuning on small CI GPUs. + # Assumes prefix caching off and max_num_seqs*max_model_len <= ~175k tokens; + # raise this cap if either grows or enable_prefix_caching is turned on. + kv_cache_memory_bytes: 6442450944 # 6 GiB enforce_eager: true trust_remote_code: true enable_prefix_caching: false diff --git a/vllm_omni/diffusion/attention/backends/sage_attn.py b/vllm_omni/diffusion/attention/backends/sage_attn.py index 080a241c1..37eab409a 100644 --- a/vllm_omni/diffusion/attention/backends/sage_attn.py +++ b/vllm_omni/diffusion/attention/backends/sage_attn.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + import torch from vllm.logger import init_logger @@ -9,17 +11,43 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionImpl, AttentionMetadata, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) -try: - from sageattention import sageattn -except ImportError: - logger.warning( - "SageAttentionBackend is not available. You may install sage-attention" - " by pip install git+https://github.com/thu-ml/SageAttention.git" - ) - raise ImportError +_SAGE_ATTN_TENSOR_LAYOUT = os.environ.get("SAGE_ATTN_TENSOR_LAYOUT", "NHD").upper() +assert _SAGE_ATTN_TENSOR_LAYOUT in ("NHD", "HND"), ( + f"SAGE_ATTN_TENSOR_LAYOUT must be 'NHD' or 'HND', got '{_SAGE_ATTN_TENSOR_LAYOUT}'" +) + +if current_omni_platform.is_xpu(): + try: + import inspect + + from auto_round_kernel import ARK + + _ark = ARK() + xpu_sageattn = _ark.sagev1 + _sagev1_params = inspect.signature(xpu_sageattn).parameters + _sagev1_has_tensor_layout = "tensor_layout" in _sagev1_params + _sagev1_scale_param = "sm_scale" if "sm_scale" in _sagev1_params else "scale" + except ImportError: + logger.warning( + "XPU SageAttention (auto_round_kernel.ARK.sagev1) is not available. " + "Install auto-round-lib for XPU sage attention support." + ) + xpu_sageattn = None + _sagev1_has_tensor_layout = False + _sagev1_scale_param = "scale" +else: + try: + from sageattention import sageattn + except ImportError: + logger.warning( + "SageAttentionBackend is not available. You may install sage-attention" + " by pip install git+https://github.com/thu-ml/SageAttention.git" + ) + raise ImportError # TODO add sage3 attention backend @@ -73,3 +101,50 @@ class SageAttentionImpl(AttentionImpl): sm_scale=self.softmax_scale, ) return output + + def forward_xpu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata = None, + ) -> torch.Tensor: + if xpu_sageattn is None: + raise ImportError("XPU SageAttention requires auto-round-lib. Install with: pip install auto-round-lib") + orig_dtype = query.dtype + q = query.to(torch.float16) if orig_dtype != torch.float16 else query + k = key.to(torch.float16) if orig_dtype != torch.float16 else key + v = value.to(torch.float16) if orig_dtype != torch.float16 else value + + if _sagev1_has_tensor_layout: + if _SAGE_ATTN_TENSOR_LAYOUT == "HND": + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + output = xpu_sageattn( + q, + k, + v, + tensor_layout=_SAGE_ATTN_TENSOR_LAYOUT, + is_causal=self.causal, + **{_sagev1_scale_param: self.softmax_scale}, + ) + if _SAGE_ATTN_TENSOR_LAYOUT == "HND": + output = output.transpose(1, 2).contiguous() + else: + # No tensor_layout support: kernel expects HND [B, H, S, D] + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + output = xpu_sageattn( + q, + k, + v, + is_causal=self.causal, + **{_sagev1_scale_param: self.softmax_scale}, + ) + output = output.transpose(1, 2).contiguous() + + if orig_dtype != torch.float16: + output = output.to(orig_dtype) + return output diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 6a2600e9f..ff1eea6ae 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -242,6 +242,80 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], return refresh_cache_context +def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for Wan2.2 S2V. + + S2V uses a single transformer, but unlike the other Wan2.2 variants its + block loop calls each block as ``block(hidden_states, **kwargs)`` and keeps + the timestep modulation state in ``e`` rather than a second positional + tensor. CacheDiT Pattern_3 matches that contract: cache hidden states only + and pass the remaining conditioning through kwargs unchanged. + + The S2V transformer has an ``after_transformer_block`` method that injects + audio embeddings after specific layers. The cached blocks wrapper + (Wan22S2VCachedBlocks._run_block) calls the original internally, so we + permanently replace it with a no-op on the transformer to prevent double + injection from the main forward loop. + """ + db_cache_config = _build_db_cache_config(cache_config) + calibrator_config = None + if cache_config.enable_taylorseer: + taylorseer_order = cache_config.taylorseer_order + calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) + logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + + # Save the original after_transformer_block before cache-dit wrapping + transformer = pipeline.transformer + if hasattr(transformer, "after_transformer_block"): + transformer._cache_dit_original_after_transformer_block = transformer.after_transformer_block + + Wan22S2VCachedAdapter.apply( + BlockAdapter( + transformer=transformer, + blocks=[transformer.blocks], + forward_pattern=[ForwardPattern.Pattern_3], + params_modifiers=[ + ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config), + ], + has_separate_cfg=True, + ), + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + + # Permanently replace after_transformer_block with a no-op. + # The cached blocks wrapper (Wan22S2VCachedBlocks._run_block) already calls + # the original via _cache_dit_original_after_transformer_block. + def _noop_after_transformer_block(block_idx: int, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + transformer.after_transformer_block = _noop_after_transformer_block + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context for the S2V transformer.""" + if cache_config.scm_steps_mask_policy is None: + cache_dit.refresh_context( + pipeline.transformer, + num_inference_steps=num_inference_steps, + verbose=verbose, + ) + else: + cache_dit.refresh_context( + pipeline.transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_inference_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, + total_steps=num_inference_steps, + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + + def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]: """Enable cache-dit for LongCatImage pipeline. @@ -986,6 +1060,151 @@ class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2): return hidden_states, encoder_hidden_states +class Wan22S2VCachedBlocks(CachedBlocks_Pattern_3_4_5): + """CacheDiT blocks wrapper that preserves S2V per-layer audio injection.""" + + def _run_block(self, block_id: int, block: torch.nn.Module, hidden_states: torch.Tensor, *args, **kwargs): + hidden_states = block(hidden_states, *args, **kwargs) + hidden_states, new_encoder_hidden_states = self._process_block_outputs(hidden_states) + original_after_transformer_block = getattr( + self.transformer, + "_cache_dit_original_after_transformer_block", + getattr(self.transformer, "after_transformer_block", None), + ) + if original_after_transformer_block is not None: + hidden_states = original_after_transformer_block(block_id, hidden_states) + return hidden_states, new_encoder_hidden_states + + def call_blocks( + self, + hidden_states: torch.Tensor, + *args, + **kwargs, + ): + new_encoder_hidden_states = None + for block_id, block in enumerate(self.transformer_blocks): + hidden_states, new_encoder_hidden_states = self._run_block( + block_id, + block, + hidden_states, + *args, + **kwargs, + ) + return hidden_states, new_encoder_hidden_states + + def call_Fn_blocks( + self, + hidden_states: torch.Tensor, + *args, + **kwargs, + ): + new_encoder_hidden_states = None + for block_id, block in enumerate(self._Fn_blocks()): + hidden_states, new_encoder_hidden_states = self._run_block( + block_id, + block, + hidden_states, + *args, + **kwargs, + ) + return hidden_states, new_encoder_hidden_states + + def call_Mn_blocks( + self, + hidden_states: torch.Tensor, + *args, + **kwargs, + ): + original_hidden_states = hidden_states + new_encoder_hidden_states = None + start_idx = self.context_manager.Fn_compute_blocks() + for block_id, block in enumerate(self._Mn_blocks(), start=start_idx): + hidden_states, new_encoder_hidden_states = self._run_block( + block_id, + block, + hidden_states, + *args, + **kwargs, + ) + + hidden_states = hidden_states.contiguous() + hidden_states_residual = hidden_states - original_hidden_states.to(hidden_states.device) + + return ( + hidden_states, + new_encoder_hidden_states, + hidden_states_residual, + ) + + def call_Bn_blocks( + self, + hidden_states: torch.Tensor, + *args, + **kwargs, + ): + new_encoder_hidden_states = None + if self.context_manager.Bn_compute_blocks() == 0: + return hidden_states, new_encoder_hidden_states + + start_idx = len(self.transformer_blocks) - self.context_manager.Bn_compute_blocks() + for block_id, block in enumerate(self._Bn_blocks(), start=start_idx): + hidden_states, new_encoder_hidden_states = self._run_block( + block_id, + block, + hidden_states, + *args, + **kwargs, + ) + + return hidden_states, new_encoder_hidden_states + + +class Wan22S2VCachedAdapter(CachedAdapter): + """CacheDiT adapter that uses Wan22S2VCachedBlocks for S2V audio injection. + + Only overrides collect_unified_blocks to use Wan22S2VCachedBlocks (which + calls after_transformer_block per-layer internally). The base class + mock_transformer handles the forward wrapping — after_transformer_block is + permanently replaced with a no-op in enable_cache_for_wan22_s2v() to prevent + double injection. + """ + + @classmethod + def collect_unified_blocks( + cls, + block_adapter: BlockAdapter, + contexts_kwargs: list[dict], + ) -> list[dict[str, torch.nn.ModuleList]]: + BlockAdapter.assert_normalized(block_adapter) + + total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = [] + assert hasattr(block_adapter.pipe, "_context_manager") + + for i in range(len(block_adapter.transformer)): + unified_blocks_bind_context = {} + for j in range(len(block_adapter.blocks[i])): + cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"] + unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList( + [ + Wan22S2VCachedBlocks( + block_adapter.blocks[i][j], + transformer=block_adapter.transformer[i], + forward_pattern=block_adapter.forward_pattern[i][j], + check_forward_pattern=block_adapter.check_forward_pattern, + check_num_outputs=block_adapter.check_num_outputs, + cache_prefix=block_adapter.blocks_name[i][j], + cache_context=block_adapter.unique_blocks_name[i][j], + context_manager=block_adapter.pipe._context_manager, + cache_type=cache_config.cache_type, + ) + ] + ) + + total_cached_blocks.append(unified_blocks_bind_context) + + return total_cached_blocks + + class BagelCachedAdapter(CachedAdapter): """ Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks. @@ -1808,6 +2027,7 @@ CUSTOM_DIT_ENABLERS.update( { "Wan22Pipeline": enable_cache_for_wan22, "Wan22I2VPipeline": enable_cache_for_wan22, + "Wan22S2VPipeline": enable_cache_for_wan22_s2v, "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3, "FluxPipeline": enable_cache_for_flux, "Flux2KleinPipeline": enable_cache_for_flux2_klein, diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py index f808b301b..60638cca8 100644 --- a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py +++ b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py @@ -81,7 +81,7 @@ class DefaultAdapter: loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config) # load_model will handle dtypes / device placement, put in .eval() mode - return loader.load_model(od_config=od_config, load_device=device) + return loader.load_model(load_device=device) @staticmethod def get_transformer(pipeline: Any) -> tuple[Any, str]: diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 5744c6faa..4017ac984 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.base_config import ( from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.network_utils import is_port_available +from vllm_omni.errors import client_error_metadata from vllm_omni.quantization import build_quant_config if TYPE_CHECKING: @@ -610,9 +611,6 @@ class OmniDiffusionConfig: # Step mode settings step_execution: bool = False - # sleep mode - enable_sleep_mode: bool = False - # Maximum number of sequences to generate in a batch max_num_seqs: int = 1 @@ -1065,6 +1063,8 @@ class DiffusionOutput: trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None trajectory_decoded: list[Image.Image] | None = None error: str | None = None + error_status_code: int | None = None + error_type: str | None = None aborted: bool = False abort_message: str | None = None @@ -1105,6 +1105,15 @@ class DiffusionOutput: if self.custom_output: self.custom_output = {k: _maybe_to_cpu(v) for k, v in self.custom_output.items()} + @classmethod + def from_exception(cls, exc: BaseException) -> "DiffusionOutput": + status_code, error_type = client_error_metadata(exc) + return cls( + error=str(exc), + error_status_code=status_code, + error_type=error_type, + ) + class DiffusionRequestAbortedError(RuntimeError): """Raised when a diffusion request ends via user-visible abort.""" diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 54e93c718..ce1185c66 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -34,6 +34,7 @@ from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusio from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput +from vllm_omni.errors import client_error_from_metadata, is_client_error_status from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.outputs import OmniRequestOutput @@ -96,6 +97,8 @@ def _move_tensor_tree_to_cpu(value: object) -> object: def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int: + """Get num_frames for the dummy warmup run. Returns 0 to skip warmup.""" + model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"): return int(getattr(model_cls, "dummy_run_num_frames")) @@ -229,6 +232,12 @@ class DiffusionEngine: if output.aborted: raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.") if output.error: + if is_client_error_status(output.error_status_code): + raise client_error_from_metadata( + output.error, + status_code=output.error_status_code, + error_type=output.error_type, + ) raise RuntimeError(output.error) logger.debug("Generation completed successfully.") @@ -505,7 +514,7 @@ class DiffusionEngine: request_id=request_id, step_index=None, finished=True, - result=DiffusionOutput(error=str(exc)), + result=DiffusionOutput.from_exception(exc), ) for request_id in sched_output.scheduled_request_ids ] @@ -662,7 +671,7 @@ class DiffusionEngine: request_id=request_id, step_index=None, finished=True, - result=DiffusionOutput(error=str(exc)), + result=DiffusionOutput.from_exception(exc), ) self._process_aborts_queue() @@ -722,6 +731,9 @@ class DiffusionEngine: prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) + if num_frames <= 0: + logger.info("Skipping dummy warmup run (num_frames=0)") + return req = OmniDiffusionRequest( prompts=[prompt], request_id=DUMMY_DIFFUSION_REQUEST_ID, diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py new file mode 100644 index 000000000..0e7410808 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +from diffusers import AutoencoderKLLTX2Video +from diffusers.models.autoencoders.vae import DecoderOutput +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeExecutor, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +logger = init_logger(__name__) + + +class LTX2VaeExecutor(DistributedVaeExecutor): + def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype): + tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype) + for idx, (_, t_tensor) in enumerate(local_results): + slices = tuple(slice(0, s) for s in t_tensor.shape) + tile_tensor[idx][slices] = t_tensor + return tile_tensor + + def _unpack_tiles_with_known_metadata(self, tile_gather, assigned, tile_output_shapes): + coord_tensor_map = {} + for r, rank_tasks in enumerate(assigned): + tiles_src = tile_gather[r] + for idx, task in enumerate(rank_tasks): + output_shape = tile_output_shapes[task.tile_id] + slices = tuple(slice(0, dim) for dim in output_shape) + coord_tensor_map[task.grid_coord] = tiles_src[idx][slices] + + return coord_tensor_map + + def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True): + pp_size = min(self.parallel_size, self.world_size) + + tiletask_list, grid_spec = operator.split(z) + max_tile_output_shape = grid_spec.tile_spec.get("max_tile_output_shape") + tile_output_shapes = grid_spec.tile_spec.get("tile_output_shapes") + if max_tile_output_shape is None or tile_output_shapes is None: + return super().execute(z, operator, broadcast_result=broadcast_result) + + assigned = self._balance_tasks(tiletask_list, pp_size) + local_tasks = assigned[self.rank] if self.rank < pp_size else [] + local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks] + + output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype + max_local_tile_count = max(len(tasks) for tasks in assigned) + global_padding_shape = [max_local_tile_count, *max_tile_output_shape] + local_tile_tensor = self._pack_local_tiles_without_meta( + local_results, + global_padding_shape, + z.device, + output_dtype, + ) + tile_gather = self.gather_tensors(local_tile_tensor) + + if self.rank != 0: + result = torch.empty(0, device=z.device) + else: + coord_tensor_map = self._unpack_tiles_with_known_metadata(tile_gather, assigned, tile_output_shapes) + result = operator.merge(coord_tensor_map, grid_spec) + + if broadcast_result: + result = self._sync_final_result(result, z.ndim, z.device, output_dtype) + return result + + +class DistributedAutoencoderKLLTX2Video(AutoencoderKLLTX2Video, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def init_distributed(self): + self.distributed_executor = LTX2VaeExecutor() + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + sample_num_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 + output_channels = int(getattr(getattr(self, "config", None), "out_channels", 3)) + + tiletask_list = [] + tile_output_shapes = {} + for i in range(0, height, tile_latent_stride_height): + for j in range(0, width, tile_latent_stride_width): + tile = z[:, :, :num_frames, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + tile_id = len(tiletask_list) + tile_output_shapes[tile_id] = ( + z.shape[0], + output_channels, + sample_num_frames, + tile.shape[3] * self.spatial_compression_ratio, + tile.shape[4] * self.spatial_compression_ratio, + ) + tiletask_list.append( + TileTask( + tile_id, + (i // tile_latent_stride_height, j // tile_latent_stride_width), + tile, + workload=tile.shape[2] * tile.shape[3] * tile.shape[4], + ) + ) + + tile_spec = { + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": self.tile_sample_min_height - self.tile_sample_stride_height, + "blend_width": self.tile_sample_min_width - self.tile_sample_stride_width, + "tile_sample_stride_height": self.tile_sample_stride_height, + "tile_sample_stride_width": self.tile_sample_stride_width, + "max_tile_output_shape": ( + z.shape[0], + output_channels, + sample_num_frames, + self.tile_sample_min_height, + self.tile_sample_min_width, + ), + "tile_output_shapes": tile_output_shapes, + } + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def tile_exec( + self, + task: TileTask, + temb: torch.Tensor | None = None, + causal: bool | None = None, + ) -> torch.Tensor: + if hasattr(self, "clear_cache"): + self.clear_cache() + return self.decoder(task.tensor, temb, causal=causal) + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + + if hasattr(self, "clear_cache"): + self.clear_cache() + + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append( + tile[ + :, + :, + :, + : grid_spec.tile_spec["tile_sample_stride_height"], + : grid_spec.tile_spec["tile_sample_stride_width"], + ] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=3)[ + :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"] + ] + return dec + + def tiled_decode( + self, + z: torch.Tensor, + temb: torch.Tensor | None = None, + causal: bool | None = None, + return_dict: bool = True, + ): + if not self.is_distributed_enabled(): + if causal is None: + return super().tiled_decode(z, temb, return_dict=return_dict) + return super().tiled_decode(z, temb, causal=causal, return_dict=return_dict) + + logger.debug("LTX2 video VAE decode running with distributed tiled executor") + result = self.distributed_executor.execute( + z, + DistributedOperator( + split=self.tile_split, + exec=lambda task: self.tile_exec(task, temb=temb, causal=causal), + merge=self.tile_merge, + ), + broadcast_result=False, + ) + if not return_dict: + return (result,) + + return DecoderOutput(sample=result) diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index a2717bcc8..137f2b887 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -22,6 +22,7 @@ from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.engine.stage_client import StageClientBase from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -143,11 +144,13 @@ class InlineStageDiffusionClient(StageClientBase): logger.info("request_id: %s aborted: %s", request_id, str(e)) except Exception as e: logger.exception("Diffusion request %s failed: %s", request_id, e) - error_output = OmniRequestOutput.from_diffusion( + status_code, error_type = client_error_metadata(e) + error_output = OmniRequestOutput.from_error( request_id=request_id, - images=[], + error_message=str(e), + status_code=status_code, + error_type=error_type, ) - error_output.error = str(e) self._output_queue.put_nowait(error_output) finally: self._tasks.pop(request_id, None) @@ -248,11 +251,13 @@ class InlineStageDiffusionClient(StageClientBase): logger.info("request_id: %s aborted: %s", request_id, str(e)) except Exception as e: logger.exception("Batch diffusion request %s failed: %s", request_id, e) - error_output = OmniRequestOutput.from_diffusion( + status_code, error_type = client_error_metadata(e) + error_output = OmniRequestOutput.from_error( request_id=request_id, - images=[], + error_message=str(e), + status_code=status_code, + error_type=error_type, ) - error_output.error = str(e) self._output_queue.put_nowait(error_output) finally: self._tasks.pop(request_id, None) diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index b4fa1209f..4cb714465 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -5,6 +5,7 @@ from einops import rearrange, repeat from vllm.logger import init_logger from vllm_omni.diffusion.layers.custom_op import CustomOp +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -300,3 +301,234 @@ def apply_rope_to_qk( query = rope(query, cos, sin) key = rope(key, cos, sin) return query, key + + +class WanS2VRotaryPosEmbed(torch.nn.Module): + """Precompute complex-valued RoPE embeddings for S2V multi-grid positions. + + Owns the base frequency buffer and provides forward() to compute position + embeddings given hidden_states (for shape) and grid_sizes. + """ + + def __init__(self, num_heads: int, head_dim: int, max_seq_len: int = 1024): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + d = head_dim + freqs = torch.cat( + [ + self._rope_params(max_seq_len, d - 4 * (d // 6)), + self._rope_params(max_seq_len, 2 * (d // 6)), + self._rope_params(max_seq_len, 2 * (d // 6)), + ], + dim=1, + ) + self.register_buffer("freqs", freqs.to(torch.complex64), persistent=False) + + @staticmethod + @torch.amp.autocast(current_omni_platform.device_type, enabled=False) + def _rope_params(max_seq_len, dim, theta=10000): + if dim % 2 != 0: + raise ValueError(f"dim ({dim}) must be even") + freqs = torch.outer( + torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)) + ) + return torch.polar(torch.ones_like(freqs), freqs) + + def forward( + self, + hidden_states: torch.Tensor, + grid_sizes: list, + trainable_freqs: torch.Tensor | None = None, + ) -> torch.Tensor: + """Precompute RoPE embeddings for the given grid layout. + + Args: + hidden_states: Tensor [B, S, ...] (used for batch/seq shape and device) + grid_sizes: Grid specification (list of [offsets, sizes, totals]) + trainable_freqs: Optional trainable frequency overrides for t_f < 0 + + Returns: + Complex tensor [B, S, 1, head_dim//2] of precomputed position embeddings + """ + b, s = hidden_states.shape[0], hidden_states.shape[1] + c = self.head_dim // 2 + device = hidden_states.device + + freqs = self.freqs.to(device) + if trainable_freqs is not None: + freqs_input = [freqs, trainable_freqs] + else: + freqs_input = freqs + + if isinstance(freqs_input, list): + trainable_f = freqs_input[1] + freqs_split = freqs_input[0] + else: + trainable_f = None + freqs_split = freqs_input + freqs_split = freqs_split.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + output = torch.empty((b, s, 1, c), device=device, dtype=torch.complex64) + seq_bucket = [0] + if not isinstance(grid_sizes, list): + grid_sizes = [grid_sizes] + for g in grid_sizes: + if not isinstance(g, list): + g = [torch.zeros_like(g), g] + batch_size = g[0].shape[0] + for i in range(batch_size): + f_o, h_o, w_o = g[0][i] + f, h, w = g[1][i] + t_f, t_h, t_w = g[2][i] + seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o + seq_len = int(seq_f * seq_h * seq_w) + if seq_len > 0: + if t_f > 0: + assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0 + seq_f_int = int(seq_f) + seq_h_int = int(seq_h) + seq_w_int = int(seq_w) + + if f_o >= 0: + f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long() + else: + f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long() + h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long() + w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long() + + freqs_0 = freqs_split[0][f_sam] if f_o >= 0 else freqs_split[0][f_sam].conj() + freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1) + + freqs_i = torch.cat( + [ + freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1), + freqs_split[1][h_sam] + .view(1, seq_h_int, 1, -1) + .expand(seq_f_int, seq_h_int, seq_w_int, -1), + freqs_split[2][w_sam] + .view(1, 1, seq_w_int, -1) + .expand(seq_f_int, seq_h_int, seq_w_int, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + elif t_f < 0: + freqs_i = trainable_f.unsqueeze(1) + output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i + seq_bucket.append(seq_bucket[-1] + seq_len) + return output + + +class RotaryEmbeddingWanS2V(RotaryEmbeddingWan): + """Apply RoPE using precomputed complex freqs for Wan S2V main transformer. + + Converts complex freqs (from WanS2VRotaryPosEmbed) to cos/sin and delegates + to RotaryEmbeddingWan for platform-optimized application (float32 kernel). + Under TP, freqs has 1 head — broadcasts automatically via cos/sin. + """ + + def __init__(self) -> None: + super().__init__(is_neox_style=False, half_head_dim=True) + + def forward(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + freqs_sliced = freqs[:, : x.size(1)] + cos = freqs_sliced.real.to(x.dtype) + sin = freqs_sliced.imag.to(x.dtype) + return super().forward(x, cos, sin) + + +class RotaryEmbeddingS2VGrid(torch.nn.Module): + """Grid-based RoPE for S2V motioner/init attention. + + Applies complex-valued rotary embeddings using 3D grid sampling + (frame, height, width). Used by SimpleSelfAttention, SwinSelfAttention, + CausalSelfAttention in motioner blocks. + """ + + @staticmethod + @torch.amp.autocast(current_omni_platform.device_type, enabled=False) + def precompute( + seq_len: int, num_heads: int, head_dim: int, grid_sizes, freqs: torch.Tensor, device: torch.device, start=None + ) -> torch.Tensor: + """Precompute position frequency tensor from grid specification. + + Returns a complex tensor that can be reused across layers via apply_precomputed(). + """ + c = head_dim // 2 + + trainable_freqs = None + if isinstance(freqs, list): + trainable_freqs = freqs[1] + freqs = freqs[0] + freqs = freqs.to(device) + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + if not isinstance(grid_sizes, list): + grid_sizes = [grid_sizes] + + batch_size = grid_sizes[0][0].shape[0] if isinstance(grid_sizes[0], list) else grid_sizes[0].shape[0] + precomputed = torch.empty((batch_size, seq_len, 1, c), device=device, dtype=torch.complex64) + seq_bucket = [0] + + for g in grid_sizes: + if not isinstance(g, list): + g = [torch.zeros_like(g), g] + g_batch_size = g[0].shape[0] + for i in range(g_batch_size): + if start is None: + f_o, h_o, w_o = g[0][i] + else: + f_o, h_o, w_o = start[i] + + f, h, w = g[1][i] + t_f, t_h, t_w = g[2][i] + seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o + seg_len = int(seq_f * seq_h * seq_w) + if seg_len > 0: + if t_f > 0: + seq_f_int = int(seq_f) + seq_h_int = int(seq_h) + seq_w_int = int(seq_w) + + if f_o >= 0: + f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long() + else: + f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long() + h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long() + w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long() + + freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj() + freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1) + + freqs_i = torch.cat( + [ + freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1), + freqs[1][h_sam].view(1, seq_h_int, 1, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1), + freqs[2][w_sam].view(1, 1, seq_w_int, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1), + ], + dim=-1, + ).reshape(seg_len, 1, -1) + elif t_f < 0: + freqs_i = trainable_freqs.unsqueeze(1) + precomputed[i, seq_bucket[-1] : seq_bucket[-1] + seg_len] = freqs_i + seq_bucket.append(seq_bucket[-1] + seg_len) + return precomputed + + @staticmethod + @torch.amp.autocast(current_omni_platform.device_type, enabled=False) + def apply_precomputed(x: torch.Tensor, precomputed_freqs: torch.Tensor) -> torch.Tensor: + """Apply precomputed position frequencies to input tensor.""" + n = x.size(2) + input_dtype = x.dtype + seq_len = x.size(1) + precomputed_freqs = precomputed_freqs[:, :seq_len] + x_c = torch.view_as_complex(x.to(torch.float64).reshape(x.size(0), seq_len, n, -1, 2)) + x_c = torch.view_as_real(x_c * precomputed_freqs).flatten(3) + return x_c.reshape(x.size(0), seq_len, n, -1).to(input_dtype) + + @staticmethod + @torch.amp.autocast(current_omni_platform.device_type, enabled=False) + def forward(x: torch.Tensor, grid_sizes, freqs: torch.Tensor, start=None) -> torch.Tensor: + n, d = x.size(2), x.size(3) + precomputed = RotaryEmbeddingS2VGrid.precompute(x.size(1), n, d, grid_sizes, freqs, x.device, start=start) + return RotaryEmbeddingS2VGrid.apply_precomputed(x, precomputed) diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 63e8d9a96..8ba30b79e 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -366,7 +366,11 @@ class DiffusionLoRAManager: fully_sharded_loras=False, ) - for component_name in ("transformer", "transformer_2", "dit", "bagel"): + # Default denoising components plus any a pipeline opts into via + # ``_lora_components`` (opt-in; other pipelines unchanged). + default_components = ("transformer", "transformer_2", "dit", "bagel") + extra_components = tuple(getattr(self.pipeline, "_lora_components", ()) or ()) + for component_name in (*default_components, *extra_components): if not hasattr(self.pipeline, component_name): continue component = getattr(self.pipeline, component_name) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index c9f81b64d..e5ef5e319 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import dataclasses import glob import os @@ -30,8 +31,9 @@ from vllm.transformers_utils.repo_utils import file_exists from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.torch_utils import set_default_torch_dtype +from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import OmniDiffusionConfig -from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig +from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp_to_model from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) @@ -100,18 +102,20 @@ class DiffusersPipelineLoader: counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 - def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig | None = None): + def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig): self.load_config = load_config self.od_config = od_config + self.quant_config = od_config.quantization_config + self.parallel_config = od_config.parallel_config def _prepare_weights( self, - model_name_or_path: Path, + model_name_or_path: Path | str, subfolder: str | None, revision: str | None, fall_back_to_pt: bool, allow_patterns_overrides: list[str] | None, - ) -> tuple[str, list[str], bool]: + ) -> tuple[Path | str, list[str], bool]: """Prepare weights for the model. If the model is not local, it will be downloaded.""" @@ -208,15 +212,13 @@ class DiffusersPipelineLoader: source.allow_patterns_overrides, ) - od_config = self.od_config use_multithread = ( use_safetensors - and od_config is not None - and getattr(od_config, "enable_multithread_weight_load", False) + and getattr(self.od_config, "enable_multithread_weight_load", False) and self.load_config.safetensors_load_strategy != "torchao" ) if use_multithread: - num_threads = getattr(od_config, "num_weight_load_threads", 4) + num_threads = getattr(self.od_config, "num_weight_load_threads", 4) # Keep deterministic shard order before passing to vLLM helper. sorted_hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key) weights_iterator = multi_thread_safetensors_weights_iterator( @@ -242,10 +244,7 @@ class DiffusersPipelineLoader: return prefixed_weights_iterator def _get_source_quant_config(self, source: "ComponentSource") -> object | None: - if self.od_config is None: - return None - - quant_config = self.od_config.quantization_config + quant_config = self.quant_config if hasattr(quant_config, "resolve"): return quant_config.resolve(source.prefix.rstrip(".")) return quant_config @@ -308,7 +307,6 @@ class DiffusersPipelineLoader: def load_model( self, - od_config: OmniDiffusionConfig, load_device: str, load_format: str | None = "default", custom_pipeline_name: str | type[nn.Module] | None = None, @@ -317,14 +315,13 @@ class DiffusersPipelineLoader: """Load a model with the given configurations.""" if load_format is None: load_format = "default" - self.od_config = od_config # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8), # weights are already quantized in the checkpoint — load directly on CPU. # For online quantization, load on device so quantization can run on accelerator, # then move back to CPU afterward. offload_after_quant = False - if load_device == "cpu" and od_config.quantization_config is not None: - quant_cfg = od_config.quantization_config + if load_device == "cpu" and self.quant_config is not None and device is not None: + quant_cfg = self.quant_config is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr( quant_cfg, "is_checkpoint_quantized", False ) @@ -339,48 +336,21 @@ class DiffusersPipelineLoader: logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU") target_device = torch.device(load_device) - with set_default_torch_dtype(od_config.dtype): - if od_config.parallel_config.use_hsdp: + with set_default_torch_dtype(self.od_config.dtype): + if self.parallel_config.use_hsdp: model = self._load_model_with_hsdp( - od_config, target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name + target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name ) else: - if load_format == "custom_pipeline": - # NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)` - # internally. If we construct them under `with target_device:` (CUDA), - # safetensors takes a direct-to-GPU fast path that calls `cudaMalloc` - # via the driver API and BYPASSES PyTorch's caching allocator. - # That makes those bytes invisible to CuMemAllocator, so `sleep()` - # cannot offload/unmap them and GPU memory stays pinned. - # - # Fix: build the custom pipeline on CPU first (no default device - # context), then explicitly move it to the target device. The - # subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)` - # + `copy_`, which goes through the caching allocator and is fully - # tracked by CuMemAllocator. - from vllm_omni.diffusion.config import set_current_diffusion_config - - model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name) - with set_current_diffusion_config(od_config): - model = model_cls(od_config=od_config) - if target_device.type != "cpu": - model.to(target_device) - else: - with target_device: - if load_format == "default": - model = initialize_model(od_config) - elif load_format == "diffusers": - model = DiffusersAdapterPipeline(od_config=od_config, device=target_device) - else: - raise ValueError(f"Unknown load_format: {load_format}") + model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) logger.debug("Loading weights on %s ...", load_device) if load_format == "diffusers": # DiffusersAdapterPipeline.load_weights() calls # DiffusionPipeline.from_pretrained() internally; it does # not use our native customized pipeline classes. cast(DiffusersAdapterPipeline, model).load_weights() - elif self._is_gguf_quantization(od_config): - self._load_weights_with_gguf(model, od_config) + elif self._is_gguf_quantization(): + self._load_weights_with_gguf(model) else: # Quantization does not happen in `load_weights` but after it self.load_weights(model) @@ -403,7 +373,7 @@ class DiffusersPipelineLoader: """ for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) - if isinstance(quant_method, QuantizeMethodBase): + if quant_method is not None and isinstance(quant_method, QuantizeMethodBase): # Move module to target device for processing if needed module_device = next(module.parameters(), None) if module_device is not None: @@ -433,8 +403,13 @@ class DiffusersPipelineLoader: # that have loaded weights tracking currently. if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights - # NOTE: if the model is quantized, ignore not_loaded check for scale weights - weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")} + # NOTE: if the model is quantized, ignore not_loaded check for scale + # weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a + # static activation `input_scale`, which the quant method may + # fold/track differently than plain parameters. + weights_scale_not_loaded = { + name for name in weights_not_loaded if name.endswith(("weight_scale", "input_scale")) + } weights_not_loaded = weights_not_loaded - weights_scale_not_loaded if weights_not_loaded: self._check_unloaded_weights(weights_not_loaded) @@ -476,9 +451,11 @@ class DiffusersPipelineLoader: are logged as a warning. Any *other* missing weight raises ``ValueError`` regardless of quantization. """ - od_config = getattr(self, "od_config", None) - if od_config is None or od_config.quantization_config is None: - raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}") + if self.quant_config is None: + raise ValueError( + "The quantization config is None, and the following weights " + f"were not initialized from checkpoint: {weights_not_loaded}" + ) expected_missing = {w for w in weights_not_loaded if self._is_expected_quantized_weight(w)} unexpected_missing = weights_not_loaded - expected_missing @@ -491,39 +468,33 @@ class DiffusersPipelineLoader: if unexpected_missing: raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}") - def _is_gguf_quantization(self, od_config: OmniDiffusionConfig) -> bool: - quant_config = od_config.quantization_config - if quant_config is None: - return False + def _get_gguf_model_from_config(self) -> str | None: + """Extract gguf_model from quant_config (handles both dict and object).""" + if self.quant_config is None: + return None + if isinstance(self.quant_config, dict): + return self.quant_config.get("gguf_model") + return getattr(self.quant_config, "gguf_model", None) - # New API: DiffusionGGUFConfig (QuantizationConfig subclass) + def _is_gguf_quantization(self) -> bool: + """Check whether or not this pipeline loader is pointing at a GGUF config.""" from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig - if isinstance(quant_config, DiffusionGGUFConfig): - if quant_config.gguf_model is None: + maybe_gguf_model = self._get_gguf_model_from_config() + is_gguf = ( + isinstance(self.quant_config, DiffusionGGUFConfig) + or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf") + or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf") + ) + if is_gguf: + if not maybe_gguf_model: raise ValueError("GGUF quantization requires gguf_model") return True - # Dict-style config: {"method": "gguf", "gguf_model": "..."} - if isinstance(quant_config, dict): - if quant_config.get("method") == "gguf": - if not quant_config.get("gguf_model"): - raise ValueError("GGUF quantization requires gguf_model") - return True + if isinstance(self.quant_config, dict): return False - # Check by name for any config that reports as "gguf" - if hasattr(quant_config, "get_name") and quant_config.get_name() == "gguf": - gguf_model = getattr(quant_config, "gguf_model", None) - if gguf_model is None: - raise ValueError("GGUF quantization requires gguf_model") - return True - - # Fallback: object with gguf_model attribute but no get_name - if hasattr(quant_config, "gguf_model") and quant_config.gguf_model: - return True - - return False + return bool(maybe_gguf_model) def _is_transformer_source(self, source: "ComponentSource") -> bool: if source.subfolder == "transformer": @@ -566,25 +537,23 @@ class DiffusersPipelineLoader: self, source: "ComponentSource", model: nn.Module, - od_config: OmniDiffusionConfig, ) -> Generator[tuple[str, torch.Tensor], None, None]: - quant_config = od_config.quantization_config - gguf_model = getattr(quant_config, "gguf_model", None) + gguf_model = self._get_gguf_model_from_config() if gguf_model is None: raise ValueError("GGUF quantization requires quantization_config.gguf_model") - gguf_file = self._resolve_gguf_model_path(gguf_model, od_config.revision) - adapter = get_gguf_adapter(gguf_file, model, source, od_config) + gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision) + adapter = get_gguf_adapter(gguf_file, model, source, self.od_config) weights_iter = adapter.weights_iterator() return ((source.prefix + name, tensor) for (name, tensor) in weights_iter) - def _load_weights_with_gguf(self, model: nn.Module, od_config: OmniDiffusionConfig) -> set[str]: + def _load_weights_with_gguf(self, model: nn.Module) -> set[str]: sources = self._get_weight_sources(model) loaded: set[str] = set() loadable_names: set[str] | None = None for source in sources: if self._is_transformer_source(source): - loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model, od_config)) + loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model)) # GGUF checkpoints can be transformer-only or partially quantized. # Only fall back to HF if this source still has missing loadable weights. @@ -609,9 +578,45 @@ class DiffusersPipelineLoader: raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}") return loaded + def _init_from_load_format( + self, + load_format: str, + target_device: torch.device, + custom_pipeline_name: str | type[nn.Module] | None = None, + is_hsdp: bool = False, + ) -> nn.Module: + """Initialize the model from a specified load format.""" + if load_format == "custom_pipeline": + # NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)` + # internally. If we construct them under `with target_device:` (CUDA), + # safetensors takes a direct-to-GPU fast path that calls `cudaMalloc` + # via the driver API and BYPASSES PyTorch's caching allocator. + # That makes those bytes invisible to CuMemAllocator, so `sleep()` + # cannot offload/unmap them and GPU memory stays pinned. + # + # Fix: build the custom pipeline on CPU first (no default device + # context), then explicitly move it to the target device. The + # subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)` + # + `copy_`, which goes through the caching allocator and is fully + # tracked by CuMemAllocator. + model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name) + with set_current_diffusion_config(self.od_config): + model = model_cls(od_config=self.od_config) + if not is_hsdp and target_device.type != "cpu": + model.to(target_device) + else: + device_ctx = target_device if not is_hsdp else contextlib.nullcontext() + with device_ctx: + if load_format == "default": + model = initialize_model(self.od_config) + elif load_format == "diffusers": + model = DiffusersAdapterPipeline(od_config=self.od_config, device=target_device) + else: + raise ValueError(f"Unknown load_format: {load_format}") + return model + def _load_model_with_hsdp( self, - od_config: OmniDiffusionConfig, target_device: torch.device, load_format: str = "default", custom_pipeline_name: str | type[nn.Module] | None = None, @@ -624,14 +629,11 @@ class DiffusersPipelineLoader: Approach: Load weights first using model's load_weights (handles QKV fusion etc.), then apply HSDP sharding to redistribute weights across GPUs. """ - from vllm_omni.diffusion.distributed.hsdp import apply_hsdp_to_model - - parallel_config = od_config.parallel_config hsdp_config = HSDPInferenceConfig( enabled=True, - hsdp_replicate_size=parallel_config.hsdp_replicate_size, - hsdp_shard_size=parallel_config.hsdp_shard_size, - param_dtype=od_config.dtype, + hsdp_replicate_size=self.parallel_config.hsdp_replicate_size, + hsdp_shard_size=self.parallel_config.hsdp_shard_size, + param_dtype=self.od_config.dtype, ) # Initialize model WITHOUT device context (weights start on CPU). @@ -639,14 +641,9 @@ class DiffusersPipelineLoader: # directly on GPU, HSDP needs weights on CPU first so they can be redistributed # across GPUs by apply_hsdp_to_model. The model's load_weights handles weight # mapping (QKV fusion, etc.). - if load_format == "default": - model = initialize_model(od_config) - elif load_format == "custom_pipeline": - model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name) - from vllm_omni.diffusion.config import set_current_diffusion_config - - with set_current_diffusion_config(od_config): - model = model_cls(od_config=od_config) + if load_format == "diffusers": + raise ValueError("HSDP is not supported with the diffusers adapter load format") + model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) self.load_weights(model) # Collect all transformers to shard (some models have transformer_2 for MoE) diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 4231d4cc6..1cb20e5c9 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -1749,12 +1749,12 @@ class Bagel(nn.Module): trajectory_latents.append(x_t.clone()) for i, t in enumerate(timesteps.tolist()): # host floats; a 0-d tensor t would sync each step timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device) - if frame_condition_token_indexes is not None: - # Cond positions stay at t=0 (clean signal). Matches upstream - # PR #33 lance.py line 1605: - # timestep[current_vae_mse_indexes_local_in_vae] = t - # (cond positions remain at the ``torch.zeros`` init value). - timestep[frame_condition_token_indexes] = 0.0 + if frame_condition_token_indexes is not None: + # Cond positions stay at t=0 (clean signal). Matches upstream + # PR #33 lance.py line 1605: + # timestep[current_vae_mse_indexes_local_in_vae] = t + # (cond positions remain at the ``torch.zeros`` init value). + timestep[frame_condition_token_indexes] = 0.0 in_cfg_window = t > cfg_interval[0] and t <= cfg_interval[1] cfg_text_scale_ = cfg_text_scale if in_cfg_window else 1.0 cfg_img_scale_ = cfg_img_scale if in_cfg_window else 1.0 @@ -1818,12 +1818,12 @@ class Bagel(nn.Module): trajectory_latents.append(x_t.clone()) for i, t in enumerate(timesteps.tolist()): # host floats; a 0-d tensor t would sync each step timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device) - if frame_condition_token_indexes is not None: - # Cond positions stay at t=0 (clean signal). Matches upstream - # PR #33 lance.py line 1605: - # timestep[current_vae_mse_indexes_local_in_vae] = t - # (cond positions remain at the ``torch.zeros`` init value). - timestep[frame_condition_token_indexes] = 0.0 + if frame_condition_token_indexes is not None: + # Cond positions stay at t=0 (clean signal). Matches upstream + # PR #33 lance.py line 1605: + # timestep[current_vae_mse_indexes_local_in_vae] = t + # (cond positions remain at the ``torch.zeros`` init value). + timestep[frame_condition_token_indexes] = 0.0 v_t = self.forward_single_branch( x_t=x_t, timestep=timestep, diff --git a/vllm_omni/diffusion/models/cosmos3/guardrails.py b/vllm_omni/diffusion/models/cosmos3/guardrails.py index 3458c2cc3..e249d2826 100644 --- a/vllm_omni/diffusion/models/cosmos3/guardrails.py +++ b/vllm_omni/diffusion/models/cosmos3/guardrails.py @@ -20,6 +20,7 @@ import torch from vllm.logger import init_logger from vllm_omni.diffusion.models.progress_bar import _is_rank_zero +from vllm_omni.errors import GuardrailViolationError if TYPE_CHECKING: from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -67,7 +68,7 @@ def _build_text_guardrail(checker: Any) -> TextGuardrailFn: def text_guardrail(prompt: str) -> None: if not checker.check_text_safety(prompt): # CosmosSafetyChecker logs the specific reason at CRITICAL. - raise ValueError("Guardrail blocked prompt") + raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.") return text_guardrail diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 501f8ffab..df9eb835d 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Cosmos3 text/image-to-video and text-to-image pipeline for vllm-omni. +"""Cosmos3 text/image/video-to-video and text-to-image pipeline for vllm-omni. -Single pipeline class supports T2V, I2V, and T2I; the mode is selected at +Single pipeline class supports T2V, I2V, V2V, and T2I; the mode is selected at runtime by: * ``prompt["modalities"]`` contains ``"image"``: **T2I** (text-to-image). @@ -10,6 +10,8 @@ runtime by: (text-to-video). * ``multi_modal_data['image']`` present on the prompt: **I2V** (handled by :func:`get_cosmos3_pre_process_func`) +* ``multi_modal_data['video']`` present on the prompt with no action mode: + **V2V** (handled by :func:`get_cosmos3_pre_process_func`) """ @@ -40,7 +42,10 @@ from vllm_omni.diffusion.distributed.parallel_state import ( ) from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.interface import ( + ReferenceVideoDecodeSpec, + SupportImageInput, +) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -63,6 +68,17 @@ from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen logger = init_logger(__name__) +COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1) +COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first" +# Mirrors the WAN VAE's temporal compression. Authoritative value is +# ``self.vae.config.scale_factor_temporal`` at runtime; this constant exists so +# off-line / API code that runs before the pipeline is constructed can compute +# pixel-frame budgets without instantiating the VAE. +COSMOS3_VAE_TEMPORAL_COMPRESSION = 4 +COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = ( + max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1 +) +COSMOS3_V2V_DEFAULT_FLOW_SHIFT = 10.0 COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." COSMOS3_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution." @@ -91,16 +107,55 @@ COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0) COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096 +def _normalize_condition_frame_indexes_vision(value: Any) -> tuple[int, ...]: + """Normalize Cosmos3 vision-conditioning latent frame indexes.""" + if value is None: + return COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION + if isinstance(value, str): + value = [item.strip() for item in value.split(",") if item.strip()] + elif isinstance(value, int): + value = [value] + + if not isinstance(value, Iterable): + raise TypeError( + "Cosmos3 condition_frame_indexes_vision must be an int, comma-separated string, " + f"or iterable of ints; got {type(value)!r}." + ) + + indexes = tuple(sorted({int(index) for index in value})) + if not indexes: + raise ValueError("Cosmos3 condition_frame_indexes_vision must contain at least one index.") + if any(index < 0 for index in indexes): + raise ValueError(f"Cosmos3 condition_frame_indexes_vision must be non-negative, got {indexes}.") + return indexes + + +def _condition_pixel_frame_count( + condition_frame_indexes_vision: Iterable[int], + temporal_compression: int = COSMOS3_VAE_TEMPORAL_COMPRESSION, +) -> int: + return max(condition_frame_indexes_vision) * int(temporal_compression) + 1 + + +def _normalize_condition_video_keep(value: Any) -> str: + keep = str(value or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower() + if keep not in {"first", "last"}: + raise ValueError("Cosmos3 condition_video_keep must be either 'first' or 'last'.") + return keep + + # --------------------------------------------------------------------------- # Post-process function (registered in registry.py) # --------------------------------------------------------------------------- def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): - """Pre-process function for both T2V and I2V. + """Pre-process function for T2V, I2V, V2V, and action inputs. For T2V (no image in ``multi_modal_data``), the request is returned unchanged after the optional guardrails check. For I2V (image present), the conditioning image is loaded, aspect-resized + center-cropped, and stored back on the prompt as ``additional_information.preprocessed_image``. + For V2V (video present without action mode), source frames are cropped to + the target size and stored as ``additional_information.preprocessed_video``. """ from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled @@ -137,7 +192,27 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): return PIL.Image.open(value).convert("RGB") if isinstance(value, PIL.Image.Image): return value.convert("RGB") - raise TypeError(f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}.") + if isinstance(value, np.ndarray): + array = value + if array.ndim == 3 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4): + array = np.transpose(array, (1, 2, 0)) + if np.issubdtype(array.dtype, np.floating): + if array.min() < 0.0 or array.max() > 1.0: + array = np.clip(array, -1.0, 1.0) * 0.5 + 0.5 + array = (np.clip(array, 0.0, 1.0) * 255.0).round().astype(np.uint8) + return PIL.Image.fromarray(array).convert("RGB") + if isinstance(value, torch.Tensor): + tensor = value.detach().cpu() + if tensor.ndim == 3 and tensor.shape[0] in (3, 4): + tensor = tensor.permute(1, 2, 0) + if tensor.is_floating_point(): + if tensor.min().item() < 0.0 or tensor.max().item() > 1.0: + tensor = tensor.clamp(-1.0, 1.0) * 0.5 + 0.5 + tensor = (tensor.clamp(0.0, 1.0) * 255.0).round().to(torch.uint8) + return PIL.Image.fromarray(tensor.numpy()).convert("RGB") + raise TypeError( + f"Cosmos3 preprocessing expected PIL image, numpy array, torch tensor, or path, got {type(value)!r}." + ) def _resize_and_pad_action_image(image: PIL.Image.Image, target_h: int, target_w: int) -> PIL.Image.Image: scale = min(target_w / image.width, target_h / image.height, 1.0) @@ -170,6 +245,61 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): processed = [_preprocess_action_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in frames] return torch.stack(processed, dim=1).unsqueeze(0).contiguous() + def _preprocess_condition_image(image: PIL.Image.Image, target_h: int, target_w: int) -> torch.Tensor: + scale = max(target_w / image.width, target_h / image.height) + resize_w = int(np.ceil(scale * image.width)) + resize_h = int(np.ceil(scale * image.height)) + image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS) + left = (resize_w - target_w) // 2 + top = (resize_h - target_h) // 2 + image = image.crop((left, top, left + target_w, top + target_h)) + return video_processor.preprocess(image, height=target_h, width=target_w) + + def _video_payload_to_frames(video: Any) -> list[Any]: + if isinstance(video, list): + return video + if isinstance(video, torch.Tensor): + tensor = video.detach().cpu() + if tensor.ndim == 5: + if tensor.shape[0] != 1: + raise TypeError("Cosmos3 video preprocessing supports only batch size 1.") + tensor = tensor[0] + if tensor.ndim == 4 and tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4): + return [tensor[:, i] for i in range(tensor.shape[1])] + if tensor.ndim == 4 and tensor.shape[-1] in (3, 4): + return [tensor[i] for i in range(tensor.shape[0])] + if isinstance(video, np.ndarray): + array = video + if array.ndim == 5: + if array.shape[0] != 1: + raise TypeError("Cosmos3 video preprocessing supports only batch size 1.") + array = array[0] + if array.ndim == 4 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4): + return [array[:, i] for i in range(array.shape[1])] + if array.ndim == 4 and array.shape[-1] in (3, 4): + return [array[i] for i in range(array.shape[0])] + raise TypeError("Cosmos3 video input must be a non-empty list of frames or a single video tensor/array.") + + def _select_video_frames(frames: list[Any], max_frames: int, keep: str) -> list[Any]: + if not frames: + raise ValueError("Cosmos3 video input must contain at least one frame.") + if keep == "last": + return frames[-max_frames:] + return frames[:max_frames] + + def _preprocess_condition_video( + frames: list[Any], + target_h: int, + target_w: int, + max_frames: int, + keep: str, + ) -> torch.Tensor: + selected = _select_video_frames(frames, max_frames, keep) + processed = [ + _preprocess_condition_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in selected + ] + return torch.stack(processed, dim=1).unsqueeze(0).contiguous() + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: action_mode = _request_action_mode(request) if is_guardrails_enabled(od_config, request.sampling_params): @@ -183,16 +313,23 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): multi_modal_data = prompt.get("multi_modal_data", {}) or {} raw_image = multi_modal_data.get("image") raw_video = multi_modal_data.get("video") - if raw_image is None and not (action_mode is not None and raw_video is not None): + if raw_image is None and raw_video is None: continue + if raw_image is not None and raw_video is not None and action_mode is None: + raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.") if "additional_information" not in prompt: prompt["additional_information"] = {} + raw_video_frames: list[Any] | None = None + if raw_video is not None: + raw_video_frames = _video_payload_to_frames(raw_video) + if not raw_video_frames: + raise TypeError("Cosmos3 video input must be a non-empty list of PIL images or image paths.") + if raw_image is None: - if not isinstance(raw_video, list) or not raw_video: - raise TypeError("Cosmos3 action video input must be a non-empty list of PIL images or image paths.") - image = _pil_to_rgb(raw_video[0]) + assert raw_video_frames is not None # raw_image and raw_video can't both be None here + image = _pil_to_rgb(raw_video_frames[0]) else: image = _pil_to_rgb(raw_image) @@ -219,23 +356,38 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): int(target_h), int(target_w), ) - else: - scale = max(target_w / image.width, target_h / image.height) - resize_w = int(np.ceil(scale * image.width)) - resize_h = int(np.ceil(scale * image.height)) - image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS) - left = (resize_w - target_w) // 2 - top = (resize_h - target_h) // 2 - image = image.crop((left, top, left + target_w, top + target_h)) - - prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess( - image, height=target_h, width=target_w + elif raw_video is None: + prompt["additional_information"]["preprocessed_image"] = _preprocess_condition_image( + image, + int(target_h), + int(target_w), ) - if action_mode is not None and raw_video is not None: - if not isinstance(raw_video, list): - raise TypeError("Cosmos3 action video input must be a list of PIL images or image paths.") + else: + assert raw_video_frames is not None + extra = _extra_args(request) + condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision( + extra.get( + "condition_frame_indexes_vision", + prompt.get("condition_frame_indexes_vision"), + ) + ) + keep = _normalize_condition_video_keep( + extra.get("condition_video_keep", prompt.get("condition_video_keep")) + ) + max_frames = _condition_pixel_frame_count(condition_frame_indexes_vision) + prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video( + raw_video_frames, + int(target_h), + int(target_w), + max_frames, + keep, + ) + prompt["additional_information"]["condition_frame_indexes_vision"] = list( + condition_frame_indexes_vision + ) + if action_mode is not None and raw_video_frames is not None: prompt["additional_information"]["preprocessed_video"] = _preprocess_action_video( - raw_video, + raw_video_frames, int(target_h), int(target_w), ) @@ -345,13 +497,13 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig): class Cosmos3OmniDiffusersPipeline( nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin ): - """Cosmos3 text/image-to-video / text-to-image pipeline. + """Cosmos3 text/image/video-to-video / text-to-image pipeline. Architecture: Mixture-of-Transformers with Qwen3-VL backbone. - Understanding pathway: causal self-attention on text (runs once, K/V cached) - Generation pathway: cross-attention on noisy visual latents (runs each step) - Supports T2V, I2V, and T2I from the same class. Mode is selected at + Supports T2V, I2V, V2V, and T2I from the same class. Mode is selected at runtime: * **T2I** when ``prompt["modalities"]`` contains ``"image"``. Latent @@ -365,12 +517,45 @@ class Cosmos3OmniDiffusersPipeline( Frame 0 of the initial latent is set to the VAE-encoded conditioning image, frame-0 noise predictions are masked to zero, and the clean image latent is re-injected at frame 0 after each scheduler step. + * **V2V** when the request supplies a preprocessed video via + ``multi_modal_data['video']`` without an action mode. Explicit latent + frame indexes are kept clean with ``noisy_frame_mask`` and re-injected + after each scheduler step. * **T2V** otherwise (default video generation). """ support_image_input: ClassVar[bool] = True color_format: ClassVar[str] = "RGB" + @classmethod + def reference_video_decode_spec( + cls, + *, + num_frames: int | None = None, + extra_args: dict[str, Any] | None = None, + ) -> ReferenceVideoDecodeSpec: + extra_args = extra_args if isinstance(extra_args, dict) else {} + action_mode = normalize_action_mode(extra_args.get("action_mode")) + if action_mode is not None: + if num_frames is not None: + return ReferenceVideoDecodeSpec(max_frames=int(num_frames), keep="first") + action_chunk_size = extra_args.get("action_chunk_size") + if action_chunk_size is not None: + try: + max_frames = int(action_chunk_size) + 1 + except (TypeError, ValueError): + max_frames = None + if max_frames is not None and max_frames > 0: + return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first") + return ReferenceVideoDecodeSpec(max_frames=None, keep="first") + + condition_indexes = _normalize_condition_frame_indexes_vision(extra_args.get("condition_frame_indexes_vision")) + max_frames = _condition_pixel_frame_count(condition_indexes) + if num_frames is not None: + max_frames = min(max_frames, int(num_frames)) + keep = _normalize_condition_video_keep(extra_args.get("condition_video_keep")) + return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=keep) + def __init__( self, *, @@ -1188,6 +1373,70 @@ class Cosmos3OmniDiffusersPipeline( velocity_mask = 1.0 - condition_mask return latents, velocity_mask, image_latent + def _prepare_latents_v2v( + self, + video_tensor: torch.Tensor, + height: int, + width: int, + num_frames: int, + generator: torch.Generator, + condition_frame_indexes_vision: Iterable[int] | int | str | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare V2V latents with explicit clean conditioned latent frames.""" + del height, width + if video_tensor.ndim == 4: + video_tensor = video_tensor.unsqueeze(0) + if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3: + raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.") + if video_tensor.shape[2] < 1: + raise ValueError("Cosmos3 V2V video tensor must contain at least one frame.") + + C = self.transformer.latent_channel_size + T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial + W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial + indexes = _normalize_condition_frame_indexes_vision(condition_frame_indexes_vision) + out_of_range = [index for index in indexes if index >= T_lat] + if out_of_range: + raise ValueError( + "Cosmos3 condition_frame_indexes_vision contains indexes outside the latent video: " + f"indexes={indexes}, latent_frames={T_lat}." + ) + + noise = randn_tensor( + (1, C, T_lat, H_lat, W_lat), + generator=generator, + device=self.device, + dtype=self.dtype, + ) + condition_pixel_frames = _condition_pixel_frame_count(indexes, self.vae_scale_factor_temporal) + condition_video = video_tensor[:, :, :condition_pixel_frames] + if condition_video.shape[2] < condition_pixel_frames: + pad = condition_video[:, :, -1:].repeat(1, 1, condition_pixel_frames - condition_video.shape[2], 1, 1) + condition_video = torch.cat([condition_video, pad], dim=2) + + cond_prefix_latent = self._encode_video_tensor(condition_video) + expected_prefix = (1, C, max(indexes) + 1, H_lat, W_lat) + if ( + cond_prefix_latent.shape[0] != expected_prefix[0] + or cond_prefix_latent.shape[1] != expected_prefix[1] + or cond_prefix_latent.shape[2] < expected_prefix[2] + or cond_prefix_latent.shape[3:] != expected_prefix[3:] + ): + raise ValueError( + "Cosmos3 V2V condition latent shape mismatch: " + f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}." + ) + + condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype) + condition_latents = torch.zeros_like(noise) + for index in indexes: + condition_mask[:, :, index, :, :] = 1.0 + condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1] + latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise + velocity_mask = 1.0 - condition_mask + return latents, velocity_mask, condition_latents + def _prepare_latents_action_video( self, video_tensor: torch.Tensor, @@ -1567,19 +1816,20 @@ class Cosmos3OmniDiffusersPipeline( prompt = prompt_data negative_prompt = None image_tensor = None - action_video_tensor = None + video_tensor = None else: prompt = prompt_data.get("prompt", "") negative_prompt = prompt_data.get("negative_prompt") additional_info = prompt_data.get("additional_information", {}) or {} image_tensor = additional_info.get("preprocessed_image") - action_video_tensor = additional_info.get("preprocessed_video") + video_tensor = additional_info.get("preprocessed_video") sp = req.sampling_params is_t2i = self._is_t2i_request(req) sound_enabled = self._is_sound_request(prompt_data, sp) action_mode = self._get_action_mode(prompt_data, sp) action_enabled = action_mode is not None + action_video_tensor = video_tensor if action_enabled else None if action_enabled and is_t2i: raise ValueError("Cosmos3 action generation is supported only for video outputs.") if action_enabled and sound_enabled: @@ -1603,6 +1853,11 @@ class Cosmos3OmniDiffusersPipeline( ) if negative_prompt is None: negative_prompt = "" + if image_tensor is not None and video_tensor is not None and not action_enabled: + raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.") + if video_tensor is not None and is_t2i: + raise ValueError("Cosmos3 video-to-video generation is supported only for video outputs.") + is_v2v = video_tensor is not None and not is_t2i and not action_enabled # T2I and T2V share the same model + forward path; only defaults # differ: @@ -1626,7 +1881,7 @@ class Cosmos3OmniDiffusersPipeline( # Fall back to the engine-init shift, NOT None: passing None # to ``_set_flow_shift`` would leak a prior T2I rebuild # (shift=3.0) into a subsequent video request. - default_flow_shift = self._engine_init_flow_shift + default_flow_shift = COSMOS3_V2V_DEFAULT_FLOW_SHIFT if is_v2v else self._engine_init_flow_shift default_guidance_interval = None batch_size = 1 # Existing video pipeline assumes B=1. @@ -1775,6 +2030,23 @@ class Cosmos3OmniDiffusersPipeline( generator, ) image_latent = condition_latents[:, :, 0:1] + elif is_v2v: + condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision( + self._get_sp_param( + sp, + "condition_frame_indexes_vision", + self._get_prompt_param(prompt_data, "condition_frame_indexes_vision", None), + ) + ) + latents, velocity_mask, condition_latents = self._prepare_latents_v2v( + video_tensor, + height, + width, + num_frames, + generator, + condition_frame_indexes_vision, + ) + image_latent = None elif image_tensor is not None and not is_t2i: latents, velocity_mask, image_latent = self._prepare_latents_i2v( image_tensor, diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index c77777694..515e2b258 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -34,6 +34,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available from vllm_omni.diffusion.layers.norm import RMSNorm +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -1419,8 +1420,8 @@ class Cosmos3VFMTransformer(nn.Module): # For I2V: only add to noisy tokens, not conditioned ones. # Conditioned frames are clean context and should not receive # the diffusion timestep signal. - with torch.autocast("cuda", enabled=True, dtype=torch.float32): - time_embed = self.time_embedder(timestep * self.timestep_scale) + with torch.autocast(current_omni_platform.device_type, enabled=False): + time_embed = self.time_embedder((timestep * self.timestep_scale).float()) time_embed = time_embed.to(hidden_states.dtype) if noisy_frame_mask is not None: @@ -1519,7 +1520,7 @@ class Cosmos3VFMTransformer(nn.Module): if isinstance(hidden_gen, tuple): hidden_gen = hidden_gen[0] - hidden_gen = self.gen_sp_gather(hidden_gen) + hidden_gen = self.gen_sp_gather(hidden_gen) # Final norm and project back to latent space hidden_gen = self.norm_moe_gen(hidden_gen) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 80bb1d7b6..20ed54b4f 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -5,6 +5,7 @@ import inspect import logging from collections.abc import Callable, Iterable from dataclasses import dataclass +from functools import lru_cache from typing import Any, cast import numpy as np @@ -32,7 +33,6 @@ from vllm.distributed import ( get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, @@ -73,6 +73,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( ) from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx +from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding @@ -460,6 +461,22 @@ class Resolution: self.w = self.width = size[1] self.r = self.ratio = self.height / self.width + self.extra_res = set() + + def match(self, width, height) -> tuple[int, int]: + if not self.extra_res: + return self.w, self.h + else: + ret_w, ret_h = self.w, self.h + target_area = width * height + min_area_diff = abs((self.w * self.h) - target_area) + for res in self.extra_res: + area_diff = abs((res[0] * res[1]) - target_area) + if area_diff < min_area_diff: + min_area_diff = area_diff + ret_w, ret_h = res[0], res[1] + return (ret_w, ret_h) + def __getitem__(self, idx): if idx == 0: return self.h @@ -471,6 +488,19 @@ class Resolution: def __str__(self): return f"{self.h}x{self.w}" + def __repr__(self) -> str: + if not self.extra_res: + return "{" + f"{self.h}x{self.w}" + "}" + else: + ret_str = "{" + f"[{self.h}x{self.w}]" + for er in self.extra_res: + ret_str = ret_str + f"[{er[0]}x{er[1]}]" + ret_str = ret_str + "}" + return ret_str + + def append(self, res: "Resolution"): + self.extra_res.add((res.w, res.h)) + # Baked-in extras matching the official model's # `HunyuanImage3ImageProcessor.vae_reso_group` (image_processor.py:147-152). @@ -482,6 +512,10 @@ HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS: tuple[str, ...] = ( "1280x720", "768x1024", "720x1280", + "512x512", + "640x640", + "768x768", + "896x896", ) @@ -502,12 +536,17 @@ class ResolutionGroup: if extra_resolutions is not None: for er in extra_resolutions: - if not any(r.ratio == er.ratio for r in self.data): + for r in self.data: + if r.ratio == er.ratio: + r.append(er) + break + else: self.data.append(er) self.ratio = np.array([x.ratio for x in self.data]) self.attr = ["" for _ in range(len(self.data))] self.prefix_space = 0 + logger.debug(f"ResolutionGroup: {self}") def __len__(self): return len(self.data) @@ -523,13 +562,27 @@ class ResolutionGroup: res_str += ( f"\n{prefix}ID: height width ratio {' ' * max(0, attr_maxlen - 4)}count h/16 w/16 tokens\n{prefix}" ) - res_str += ("\n" + prefix).join( - [ + + rows = [] + for i, x in enumerate(self.data): + main_row = ( f"{i:2d}: ({x.h:4d}, {x.w:4d}) {self.ratio[i]:.4f} {self.attr[i]:>{attr_maxlen}s} " f"({x.h // 16:3d}, {x.w // 16:3d}) {x.h // 16 * x.w // 16:6d}" - for i, x in enumerate(self.data) - ] - ) + ) + rows.append(main_row) + extra_val = getattr(x, "extra_res", None) + if extra_val: + for sub_h, sub_w in sorted(list(extra_val)): + sub_ratio = sub_h / sub_w + sub_h16, sub_w16 = sub_h // 16, sub_w // 16 + sub_tokens = sub_h16 * sub_w16 + sub_row = ( + f" ({sub_h:4d}, {sub_w:4d}) {sub_ratio:.4f} {' ' * attr_maxlen} " + f"({sub_h16:3d}, {sub_w16:3d}) {sub_tokens:6d}" + ) + rows.append(sub_row) + + res_str += ("\n" + prefix).join(rows) res_str += f"\n{prefix_close})" return res_str @@ -568,8 +621,8 @@ class ResolutionGroup: def get_target_size(self, width, height): ratio = height / width idx = np.argmin(np.abs(self.ratio - ratio)) - reso = self.data[idx] - return reso.w, reso.h + w, h = self.data[idx].match(width, height) + return w, h def get_base_size_and_ratio_index(self, width, height): ratio = height / width @@ -577,6 +630,13 @@ class ResolutionGroup: return self.base_size, idx +@lru_cache(maxsize=4) +def get_cached_resolution_group(base_size: int) -> ResolutionGroup: + extra_res_tuple = tuple(Resolution(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS) + extra_resolutions = list(extra_res_tuple) if extra_res_tuple else None + return ResolutionGroup(base_size=base_size, extra_resolutions=extra_resolutions) + + class ImageInfo: """Class to store image information for processing and generation.""" @@ -1406,10 +1466,7 @@ class HunyuanImage3ImageProcessor: def __init__(self, config): self.config = config - self.reso_group = ResolutionGroup( - base_size=config.image_base_size, - extra_resolutions=[Resolution(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS], - ) + self.reso_group = get_cached_resolution_group(base_size=config.image_base_size) self.vae_processor = transforms.Compose( [ transforms.ToTensor(), @@ -1446,6 +1503,7 @@ class HunyuanImage3ImageProcessor: token_height = image_height // (self.config.vae_downsample_factor[0] * self.config.patch_size) token_width = image_width // (self.config.vae_downsample_factor[1] * self.config.patch_size) base_size, ratio_idx = self.reso_group.get_base_size_and_ratio_index(image_size[1], image_size[0]) + image_info = ImageInfo( image_type="gen_image", image_width=image_width, @@ -2148,7 +2206,6 @@ class HunyuanImage3Model(nn.Module): return tensor[offset * units : offset * units + units] for name, loaded_weight in weights: - # print(f"Loading weight name: {name}, tp_rank: {tp_rank}", flush=True) if contains_unexpected_keyword(name, unexpected_keywords): logger.warning("Skipping unexpected weight name: %s", name) continue diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 9cc2aafe5..5d15162fb 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -626,7 +626,7 @@ class HunyuanImage3Pipeline( rope_image_info.append(list(zip(image_slices, image_shapes))) return rope_image_info - def vae_encode(self, image, cfg_factor=1): + def vae_encode(self, image, cfg_factor=1, generator=None): config = self.vae.config if image.ndim == 3: @@ -641,11 +641,7 @@ class HunyuanImage3Pipeline( if isinstance(vae_encode_result, torch.Tensor): latents = vae_encode_result else: - # Match HunyuanImage-3's cond encode path: sample the - # posterior, but use a fixed generator so repeated online - # requests are deterministic. - _cond_vae_gen = torch.Generator(device=image.device).manual_seed(0) - latents = vae_encode_result.latent_dist.sample(_cond_vae_gen) + latents = vae_encode_result.latent_dist.sample(generator) if hasattr(config, "shift_factor") and config.shift_factor: latents.sub_(config.shift_factor) if hasattr(config, "scaling_factor") and config.scaling_factor: @@ -668,14 +664,17 @@ class HunyuanImage3Pipeline( self, batch_cond_image_info_list: list[list[JointImageInfo]], cfg_factor: int = 1, + generator=None, ): # VAE encode one by one, as we assume cond images have different sizes batch_cond_vae_images, batch_cond_t, batch_cond_vit_images = [], [], [] - for cond_image_info_list in batch_cond_image_info_list: + for batch_idx, cond_image_info_list in enumerate(batch_cond_image_info_list): cond_vae_image_list, cond_t_list, cond_vit_image_list = [], [], [] + cond_generator = generator[batch_idx] if isinstance(generator, list) else generator for image_info in cond_image_info_list: cond_t_, cond_vae_image_ = self.vae_encode( image_info.vae_image_info.image_tensor.to(self.device), + generator=cond_generator, ) cond_vit_image_list.append(image_info.vision_image_info.image_tensor) cond_vae_image_list.append(cond_vae_image_.squeeze(0)) @@ -869,7 +868,7 @@ class HunyuanImage3Pipeline( has_ar_kv = kwargs.get("ar_kv_data") if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0 and not has_ar_kv: cond_vae_images, cond_timestep, cond_vit_images = self._encode_cond_image( - batch_cond_image_info, cfg_factor[mode] + batch_cond_image_info, cfg_factor[mode], generator=generator ) vit_kwargs = {"spatial_shapes": [], "attention_mask": []} for cond_image_info in batch_cond_image_info: @@ -971,11 +970,14 @@ class HunyuanImage3Pipeline( # This implementation can handle sequences with text and image modalities, where text tokens use causal # attention and image tokens use full attention. bsz, seq_len = inputs_tensor.shape + device = inputs_tensor.device tokenizer_output = model_kwargs["tokenizer_output"] batch_image_slices = [ tokenizer_output.joint_image_slices[i] + tokenizer_output.gen_image_slices[i] for i in range(bsz) ] - attention_mask = torch.ones(seq_len, seq_len, dtype=torch.bool).tril(diagonal=0).repeat(bsz, 1, 1) + attention_mask = ( + torch.ones(seq_len, seq_len, dtype=torch.bool, device=device).tril(diagonal=0).repeat(bsz, 1, 1) + ) full_attn_spans: list[list[tuple[int, int]]] = [[] for _ in range(bsz)] for i in range(bsz): for j, image_slice in enumerate(batch_image_slices[i]): @@ -1393,18 +1395,23 @@ class HunyuanImage3Pipeline( **kwargs, ) -> DiffusionOutput: extra_args = getattr(getattr(req, "sampling_params", None), "extra_args", {}) or {} + bot_task = extra_args.get("bot_task") use_system_prompt = extra_args.get("use_system_prompt") system_prompt = extra_args.get("system_prompt") - # Fall back to per-prompt use_system_prompt forwarded by ar2diffusion + # Fall back to per-prompt prompt kwargs forwarded by ar2diffusion if req.prompts: first_prompt = req.prompts[0] if isinstance(first_prompt, dict): + if bot_task is None: + bot_task = first_prompt.get("bot_task") if use_system_prompt is None: use_system_prompt = first_prompt.get("use_system_prompt") if system_prompt is None: system_prompt = first_prompt.get("system_prompt") + if isinstance(bot_task, str) and bot_task.lower() == "none": + bot_task = None if use_system_prompt is not None: - system_prompt = get_system_prompt(use_system_prompt, "image", system_prompt) + system_prompt = get_system_prompt(use_system_prompt, bot_task or "image", system_prompt) system_prompt = system_prompt.strip() if system_prompt is not None else "" prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt @@ -1451,6 +1458,18 @@ class HunyuanImage3Pipeline( # ---- AR KV Reuse: extract injected KV from request ---- ar_kv_kwargs = self._extract_ar_kv_from_request(req) + tokenizer_bot_task = bot_task + if tokenizer_bot_task == "think_recaption": + tokenizer_bot_task = "think" + elif tokenizer_bot_task == "vanilla": + tokenizer_bot_task = "image" + supported_bot_tasks = {"auto", "image", "think", "recaption", "img_ratio"} + if tokenizer_bot_task is not None and tokenizer_bot_task not in supported_bot_tasks: + raise ValueError( + f"Unsupported HunyuanImage3 single-stage bot_task: {tokenizer_bot_task!r}. " + f"Supported values are: {sorted(supported_bot_tasks)}." + ) + model_inputs = self.prepare_model_inputs( prompt=prompt, cot_text=cot_text, @@ -1461,6 +1480,7 @@ class HunyuanImage3Pipeline( num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, batch_cond_image_info=batch_cond_image_info, + bot_task=tokenizer_bot_task or "auto", **ar_kv_kwargs, ) diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py index 1f06799c3..9af9e8488 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn.functional as F @@ -31,6 +31,9 @@ from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + logger = init_logger(__name__) @@ -332,6 +335,8 @@ class HunyuanVideo15Attention(nn.Module): out_bias: bool = True, eps: float = 1e-6, out_dim: int | None = None, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() @@ -350,6 +355,8 @@ class HunyuanVideo15Attention(nn.Module): head_size=self.head_dim, total_num_heads=self.heads, bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.to_qkv", ) self.to_out = nn.ModuleList( @@ -360,6 +367,8 @@ class HunyuanVideo15Attention(nn.Module): bias=out_bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out.0", ), nn.Identity(), # placeholder for dropout (none used) ] @@ -374,6 +383,8 @@ class HunyuanVideo15Attention(nn.Module): head_size=self.head_dim, total_num_heads=self.heads, bias=added_proj_bias, + quant_config=quant_config, + prefix=f"{prefix}.add_kv_proj", ) self.to_add_out = RowParallelLinear( @@ -382,6 +393,8 @@ class HunyuanVideo15Attention(nn.Module): bias=out_bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_add_out", ) self.rope = RotaryEmbedding(is_neox_style=False) @@ -401,6 +414,8 @@ class HunyuanVideo15Attention(nn.Module): image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, hidden_states_mask: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # Ensure contiguous for FP8 quantized linear layers + hidden_states = hidden_states.contiguous() qkv, _ = self.to_qkv(hidden_states) q_size = self.to_qkv.num_heads * self.head_dim kv_size = self.to_qkv.num_kv_heads * self.head_dim @@ -421,6 +436,7 @@ class HunyuanVideo15Attention(nn.Module): key = self.rope(key, cos, sin) if encoder_hidden_states is not None: + encoder_hidden_states = encoder_hidden_states.contiguous() encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states) add_q_size = self.add_kv_proj.num_heads * self.head_dim add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim @@ -492,6 +508,8 @@ class HunyuanVideo15TransformerBlock(nn.Module): attention_head_dim: int, mlp_ratio: float = 4.0, qk_norm: str = "rms_norm", + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ) -> None: super().__init__() hidden_size = num_attention_heads * attention_head_dim @@ -507,13 +525,23 @@ class HunyuanVideo15TransformerBlock(nn.Module): out_dim=hidden_size, bias=True, eps=1e-6, + quant_config=quant_config, + prefix=f"{prefix}.attn", ) self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio) + self.ff = FeedForward( + dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio, quant_config=quant_config, prefix=f"{prefix}.ff" + ) self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff_context = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio) + self.ff_context = FeedForward( + dim=hidden_size, + dim_out=hidden_size, + mult=mlp_ratio, + quant_config=quant_config, + prefix=f"{prefix}.ff_context", + ) def forward( self, @@ -601,6 +629,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module): target_size: int = 640, task_type: str = "i2v", use_meanflow: bool = False, + quant_config: "QuantizationConfig | None" = None, ): super().__init__() @@ -632,9 +661,14 @@ class HunyuanVideo15Transformer3DModel(nn.Module): self.transformer_blocks = nn.ModuleList( [ HunyuanVideo15TransformerBlock( - num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm + num_attention_heads, + attention_head_dim, + mlp_ratio=mlp_ratio, + qk_norm=qk_norm, + quant_config=quant_config, + prefix=f"transformer_blocks.{i}", ) - for _ in range(num_layers) + for i in range(num_layers) ] ) diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py index 038898a16..0684f6749 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py +++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py @@ -143,7 +143,9 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff self.scheduler._shift = od_config.flow_shift transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel) - self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs) + self.transformer = HunyuanVideo15Transformer3DModel( + od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs + ) # Check if model uses meanflow (distilled variants) self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False) diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py index 9eb296150..583ef87ed 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py +++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py @@ -186,7 +186,9 @@ class HunyuanVideo15I2VPipeline( self.scheduler._shift = od_config.flow_shift transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel) - self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs) + self.transformer = HunyuanVideo15Transformer3DModel( + od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs + ) self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False) diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index f0ded9cdb..d4ace761b 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -2,10 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations +from dataclasses import dataclass from typing import ( TYPE_CHECKING, Any, ClassVar, + Literal, Protocol, runtime_checkable, ) @@ -23,6 +25,12 @@ class SupportImageInput(Protocol): color_format: ClassVar[str] = "RGB" # Default color format +@dataclass(frozen=True) +class ReferenceVideoDecodeSpec: + max_frames: int | None = None + keep: Literal["first", "last"] = "first" + + @runtime_checkable class SupportAudioInput(Protocol): support_audio_input: ClassVar[bool] = True diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py index 5639e9261..6f7e9b95e 100644 --- a/vllm_omni/diffusion/models/lance/pipeline_lance.py +++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py @@ -793,17 +793,9 @@ class LancePipeline(BagelPipeline): curr_rope=cfg_text_context["ropes"], video_shapes=[video_shape], ) - cfg_img_lat = self.bagel.prepare_video_latent_cfg( - curr_kvlens=gen_context["kv_lens"], - curr_rope=gen_context["ropes"], - video_shapes=[video_shape], - ) for k, v in cfg_text_lat.items(): if torch.is_tensor(v): cfg_text_lat[k] = v.to(self.device) - for k, v in cfg_img_lat.items(): - if torch.is_tensor(v): - cfg_img_lat[k] = v.to(self.device) self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed) @@ -822,7 +814,7 @@ class LancePipeline(BagelPipeline): latents, *_ = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], cfg_text_past_key_values=cfg_text_context["past_key_values"], - cfg_img_past_key_values=gen_context["past_key_values"], # no img CFG branch + cfg_img_past_key_values=None, # no img CFG branch num_timesteps=num_timesteps, timestep_shift=timestep_shift, cfg_text_scale=cfg_text_scale, @@ -835,7 +827,7 @@ class LancePipeline(BagelPipeline): # ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens`` # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge # (derived from ``cfg_text_past_key_values``). - cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"], + cfg_img_packed_position_ids=None, # ``cfg_img_*`` index/lens kwargs removed — same as above. ) @@ -1038,11 +1030,10 @@ class LancePipeline(BagelPipeline): if torch.is_tensor(v): cfg_text_lat[k] = v.to(self.device) - # cfg_img branch: keep off (cfg_img_scale=1.0 matches upstream - # ``cfg_vit_scale=1.0``). generate_image needs the cfg_img_* args - # but skips the branch when scale<=1.0; just point them at the - # cond's lat metadata so shapes/types are fine. - cfg_img_lat = cfg_text_lat + # cfg_img branch is off by default (cfg_img_scale=1.0 matches + # upstream ``cfg_vit_scale=1.0``). Avoid passing synthetic metadata + # in that path; Bagel.generate_image only reads it when scale > 1.0. + cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed) @@ -1071,7 +1062,9 @@ class LancePipeline(BagelPipeline): # ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens`` # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge # (derived from ``cfg_text_past_key_values``). - cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"], + cfg_img_packed_position_ids=( + cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None + ), # ``cfg_img_*`` index/lens kwargs removed — same as above. ) @@ -1269,7 +1262,7 @@ class LancePipeline(BagelPipeline): for k, v in cfg_text_lat.items(): if torch.is_tensor(v): cfg_text_lat[k] = v.to(self.device) - cfg_img_lat = cfg_text_lat + cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed) @@ -1358,7 +1351,9 @@ class LancePipeline(BagelPipeline): # ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens`` # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge # (derived from ``cfg_text_past_key_values``). - cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"], + cfg_img_packed_position_ids=( + cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None + ), # ``cfg_img_*`` index/lens kwargs removed — same as above. frame_condition_token_indexes=frame_condition_token_indexes, ) @@ -1587,9 +1582,9 @@ class LancePipeline(BagelPipeline): if torch.is_tensor(v): cfg_text_lat[k] = v.to(self.device) - # cfg_img branch is off (cfg_img_scale=1.0); generate_image needs the - # arg surface but skips the branch when scale<=1.0. - cfg_img_lat = cfg_text_lat + # cfg_img branch is off by default; pass metadata only when an + # explicit cfg_img_scale enables the branch. + cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed) @@ -1618,7 +1613,9 @@ class LancePipeline(BagelPipeline): # ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens`` # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge # (derived from ``cfg_text_past_key_values``). - cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"], + cfg_img_packed_position_ids=( + cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None + ), # ``cfg_img_*`` index/lens kwargs removed — same as above. ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index a1d72cf66..4e5659c54 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -62,7 +62,17 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> kwargs["has_weight"] = elementwise_affine elif not elementwise_affine: raise TypeError("RMSNorm backend does not support disabling affine weights.") - return RMSNorm(hidden_size, **kwargs) + norm = RMSNorm(hidden_size, **kwargs) + weight = getattr(norm, "weight", None) + if ( + not elementwise_affine + and isinstance(weight, torch.Tensor) + and not isinstance(weight, nn.Parameter) + and "weight" not in dict(norm.named_buffers()) + ): + delattr(norm, "weight") + norm.register_buffer("weight", weight, persistent=False) + return norm def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index d28ff69d9..f333db8ce 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -10,7 +10,6 @@ This pipeline does NOT inherit from LTX2Pipeline because: versus LTX-2's per_layer_masked_mean_norm + shared projection path - LTX-2.3 uses a BWE vocoder outputting 48kHz audio (not 16kHz) - LTX-2.3 transformer requires the sigma parameter for prompt modulation -- CPU offloading is required for the 22B transformer (~44GB VRAM) """ from __future__ import annotations @@ -20,11 +19,11 @@ import json import os from collections.abc import Iterable from contextlib import nullcontext -from typing import Any +from typing import Any, ClassVar import numpy as np import torch -from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler +from diffusers import AutoencoderKLLTX2Audio, FlowMatchEulerDiscreteScheduler from diffusers.pipelines.ltx2 import LTX2TextConnectors from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps @@ -37,6 +36,7 @@ from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, @@ -46,7 +46,10 @@ from vllm_omni.diffusion.distributed.parallel_state import ( from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from .pipeline_ltx2 import ( @@ -112,7 +115,13 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): return post_process_func -class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): +class LTX23Pipeline( + nn.Module, + CFGParallelMixin, + ProgressBarMixin, + SupportsComponentDiscovery, + DiffusionPipelineProfilerMixin, +): """Fully independent LTX-2.3 pipeline. Key differences from LTX2Pipeline: @@ -120,11 +129,14 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): - Connectors: uses padding_side API (not additive_mask) - Vocoder: uses LTX2VocoderWithBWE (48kHz output) - Transformer: passes sigma for prompt_adaln - - CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU """ # Audio is diffused jointly with video; warmup must size audio tokens. dummy_run_num_frames = 2 + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "connectors"] + _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] + _resident_modules: ClassVar[list[str]] = ["vocoder"] def __init__( self, @@ -166,7 +178,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): # --- Tokenizer (lightweight, stays wherever) --- self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - # --- Text encoder: load on CPU, move to GPU only during encoding --- + # --- Text encoder --- with torch.device("cpu"): self.text_encoder = from_pretrained_with_prefetch( Gemma3ForConditionalGeneration.from_pretrained, @@ -177,7 +189,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): torch_dtype=dtype, ) - # --- Connectors: CPU (LTX-2.3 connectors include caption projection) --- + # --- Connectors (LTX-2.3 connectors include caption projection) --- self.connectors = from_pretrained_with_prefetch( LTX2TextConnectors.from_pretrained, model, @@ -187,9 +199,9 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): torch_dtype=dtype, ) - # --- VAE, Audio VAE: CPU --- + # --- VAE, Audio VAE --- self.vae = from_pretrained_with_prefetch( - AutoencoderKLLTX2Video.from_pretrained, + DistributedAutoencoderKLLTX2Video.from_pretrained, model, subfolder="vae", prefetch_list=ltx2_subfolders, @@ -220,6 +232,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): transformer_config = load_transformer_config(model, "transformer", local_files_only) quant_config = getattr(self.od_config, "quantization_config", None) self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) + self._place_aux_components() # --- Scheduler --- self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( @@ -261,6 +274,24 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): self._num_timesteps = None self._current_timestep = None + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + def _place_aux_components(self) -> None: + parallel_config = getattr(self.od_config, "parallel_config", None) + use_managed_placement = bool( + getattr(self.od_config, "enable_cpu_offload", False) + or getattr(self.od_config, "enable_layerwise_offload", False) + or getattr(parallel_config, "use_hsdp", False) + ) + if use_managed_placement: + return + + modules = ModuleDiscovery.discover(self) + for module in (*modules.encoders, *modules.vaes, *modules.resident_modules): + module.to(self.device) + # ------------------------------------------------------------------ # Text Encoding (LTX-2.3 specific) # ------------------------------------------------------------------ @@ -303,16 +334,11 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): text_input_ids = text_inputs.input_ids.to(device) prompt_attention_mask = text_inputs.attention_mask.to(device) - # Move text encoder to GPU for encoding - self.text_encoder.to(device) text_encoder_outputs = self.text_encoder( input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True, ) - # Move text encoder back to CPU immediately - self.text_encoder.to("cpu") - torch.accelerator.empty_cache() hidden_states = text_encoder_outputs.hidden_states @@ -958,14 +984,10 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) - self.connectors.to(device) tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left") connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side ) - self.connectors.to("cpu") - if torch.cuda.is_available(): - torch.accelerator.empty_cache() positive_connector_prompt_embeds = connector_prompt_embeds positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds @@ -1241,24 +1263,20 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): ] latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise - # Move VAE, audio_vae, vocoder to GPU for decoding - self.vae.to(device) latents = latents.to(self.vae.dtype) video = self.vae.decode(latents, timestep_decode, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - self.vae.to("cpu") + if video.numel() > 0: + video = self.video_processor.postprocess_video(video, output_type=output_type) - self.audio_vae.to(device) audio_latents = audio_latents.to(self.audio_vae.dtype) generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] - self.audio_vae.to("cpu") - self.vocoder.to(device) audio = self.vocoder(generated_mel_spectrograms) - self.vocoder.to("cpu") - torch.accelerator.empty_cache() - return DiffusionOutput(output=(video, audio)) + return DiffusionOutput( + output=(video, audio), + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index fe8695b07..fa7f95c66 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -536,6 +536,10 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli _vae_modules: ClassVar[list[str]] = [] _resident_modules: ClassVar[list[str]] = ["fm_modules"] + # Top-level module(s) the diffusion LoRA manager scans (both MoT branches). + # TODO: promote to a shared LoRA contract/mixin instead of per-pipeline opt-in. + _lora_components: ClassVar[list[str]] = ["language_model"] + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__() self.od_config = od_config @@ -1454,6 +1458,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), ] + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index 0817d6f48..1bd32bb42 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -13,7 +13,7 @@ import logging import math import os from collections.abc import Iterable -from typing import Any +from typing import Any, ClassVar import numpy as np import PIL.Image @@ -482,6 +482,7 @@ class Wan22S2VPipeline( _DEFAULT_MOTION_FRAMES = 73 _DEFAULT_INFER_FRAMES = 80 _DEFAULT_FPS = 16 + dummy_run_num_frames: ClassVar[int] = 0 def __init__( self, @@ -929,21 +930,19 @@ class Wan22S2VPipeline( def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: """Forward pass through the S2V transformer to predict noise. - The S2V model returns a list of tensors (one per batch element). - We take the first element since S2V processes one sample at a time. + With return_dict=False, the model returns (output,) where output is [B, C, T, H, W]. + We squeeze batch dim since S2V processes one sample at a time. """ if current_model is None: current_model = self.transformer - # WanModel_S2V's norm layers compute in float32; autocast ensures - # the float32→bfloat16 casts happen automatically (matching the - # original Wan2.2 inference path). param_dtype = next(current_model.parameters()).dtype with torch.amp.autocast(self.device.type, dtype=param_dtype): result = current_model(**kwargs) - # WanModel_S2V.forward returns a list of tensors if isinstance(result, list): return result[0] - return result[0] if isinstance(result, tuple) else result + if isinstance(result, tuple): + return result[0].squeeze(0) + return result.sample.squeeze(0) def diffuse( self, @@ -996,32 +995,29 @@ class Wan22S2VPipeline( latent_model_input = [latents.to(device)] timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device) - # -- Positive (conditional) prediction kwargs -- positive_kwargs = { - "x": latent_model_input, - "t": timestep, - "context": prompt_embeds[0:1], - "seq_len": max_seq_len, + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": prompt_embeds[0:1], "cond_states": cond_latents, "motion_latents": input_motion_latents, "ref_latents": ref_latents, - "motion_frames": motion_frames, "drop_motion_frames": drop_first_motion, - "audio_emb": positive_audio_emb, + "encoder_hidden_states_audio": positive_audio_emb, + "return_dict": False, } if do_true_cfg: negative_kwargs = { - "x": latent_model_input, - "t": timestep, - "context": negative_prompt_embeds[0:1], - "seq_len": max_seq_len, + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": negative_prompt_embeds[0:1], "cond_states": cond_latents, "motion_latents": input_motion_latents, "ref_latents": ref_latents, - "motion_frames": motion_frames, "drop_motion_frames": drop_first_motion, - "audio_emb": negative_audio_emb, + "encoder_hidden_states_audio": negative_audio_emb, + "return_dict": False, } else: negative_kwargs = None @@ -1174,7 +1170,7 @@ class Wan22S2VPipeline( ref_latents = self.encode_ref_image(image, height, width, device=device).to(dtype=dtype) # ---- 4. Initial motion latents (zeros) ---- - motion_frames = self.motion_frames + motion_frames = min(self.motion_frames, infer_frames - 1) motion_pixels = torch.zeros( [1, 3, motion_frames, height, width], dtype=dtype, @@ -1288,6 +1284,10 @@ class Wan22S2VPipeline( ) # ---- Decode this clip ---- + if self.od_config.enable_cpu_offload: + self.transformer.to("cpu") + current_omni_platform.empty_cache() + latents_for_decode = latents.unsqueeze(0) # [1, C, T, H, W] if not (drop_first_motion and r == 0): decode_latents = torch.cat([motion_latents, latents_for_decode], dim=2) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py index cc0bb78c1..8fef41f74 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py @@ -10,13 +10,13 @@ import math from collections.abc import Iterable from copy import deepcopy -import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.attention import AdaLayerNorm +from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps from diffusers.models.normalization import FP32LayerNorm -from einops import rearrange, repeat +from einops import rearrange from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -31,6 +31,7 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid, RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed from vllm_omni.platforms import current_omni_platform from .wan2_2_transformer import DistributedRMSNorm, WanFeedForward @@ -53,170 +54,6 @@ def sinusoidal_embedding_1d(dim, position): return x -@torch.amp.autocast(current_omni_platform.device_type, enabled=False) -def rope_params(max_seq_len, dim, theta=10000): - if dim % 2 != 0: - raise ValueError(f"dim ({dim}) must be even") - freqs = torch.outer( - torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)) - ) - freqs = torch.polar(torch.ones_like(freqs), freqs) - return freqs - - -def rope_precompute(x, grid_sizes, freqs, start=None): - """Precompute complex-valued RoPE embeddings for S2V multi-grid positions.""" - b, s, n, c = x.size(0), x.size(1), x.size(2), x.size(3) // 2 - - if isinstance(freqs, list): - trainable_freqs = freqs[1] - freqs = freqs[0] - freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) - - output = torch.view_as_complex(x.detach().reshape(b, s, n, -1, 2).to(torch.float64)) - seq_bucket = [0] - if not isinstance(grid_sizes, list): - grid_sizes = [grid_sizes] - for g in grid_sizes: - if not isinstance(g, list): - g = [torch.zeros_like(g), g] - batch_size = g[0].shape[0] - for i in range(batch_size): - if start is None: - f_o, h_o, w_o = g[0][i] - else: - f_o, h_o, w_o = start[i] - - f, h, w = g[1][i] - t_f, t_h, t_w = g[2][i] - seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o - seq_len = int(seq_f * seq_h * seq_w) - if seq_len > 0: - if t_f > 0: - if f_o >= 0: - f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1, seq_f).astype(int).tolist() - else: - f_sam = np.linspace(-f_o.item(), (-t_f - f_o).item() + 1, seq_f).astype(int).tolist() - h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1, seq_h).astype(int).tolist() - w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1, seq_w).astype(int).tolist() - - assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0 - freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj() - freqs_0 = freqs_0.view(seq_f, 1, 1, -1) - - freqs_i = torch.cat( - [ - freqs_0.expand(seq_f, seq_h, seq_w, -1), - freqs[1][h_sam].view(1, seq_h, 1, -1).expand(seq_f, seq_h, seq_w, -1), - freqs[2][w_sam].view(1, 1, seq_w, -1).expand(seq_f, seq_h, seq_w, -1), - ], - dim=-1, - ).reshape(seq_len, 1, -1) - elif t_f < 0: - freqs_i = trainable_freqs.unsqueeze(1) - output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i - seq_bucket.append(seq_bucket[-1] + seq_len) - return output - - -@torch.amp.autocast(current_omni_platform.device_type, enabled=False) -def rope_apply(x, grid_sizes, freqs, start=None): - """Apply RoPE using grid-based frequency computation (for motioner/init).""" - n, c = x.size(2), x.size(3) // 2 - input_dtype = x.dtype - - if isinstance(freqs, list): - trainable_freqs = freqs[1] - freqs = freqs[0] - freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) - - output = x.clone() - seq_bucket = [0] - if not isinstance(grid_sizes, list): - grid_sizes = [grid_sizes] - for g in grid_sizes: - if not isinstance(g, list): - g = [torch.zeros_like(g), g] - batch_size = g[0].shape[0] - for i in range(batch_size): - if start is None: - f_o, h_o, w_o = g[0][i] - else: - f_o, h_o, w_o = start[i] - - f, h, w = g[1][i] - t_f, t_h, t_w = g[2][i] - seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o - seq_len = int(seq_f * seq_h * seq_w) - if seq_len > 0: - if t_f > 0: - if f_o >= 0: - f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1, seq_f).astype(int).tolist() - else: - f_sam = np.linspace(-f_o.item(), (-t_f - f_o).item() + 1, seq_f).astype(int).tolist() - h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1, seq_h).astype(int).tolist() - w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1, seq_w).astype(int).tolist() - - assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0 - freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj() - freqs_0 = freqs_0.view(seq_f, 1, 1, -1) - - freqs_i = torch.cat( - [ - freqs_0.expand(seq_f, seq_h, seq_w, -1), - freqs[1][h_sam].view(1, seq_h, 1, -1).expand(seq_f, seq_h, seq_w, -1), - freqs[2][w_sam].view(1, 1, seq_w, -1).expand(seq_f, seq_h, seq_w, -1), - ], - dim=-1, - ).reshape(seq_len, 1, -1) - elif t_f < 0: - freqs_i = trainable_freqs.unsqueeze(1) - x_i = torch.view_as_complex( - x[i, seq_bucket[-1] : seq_bucket[-1] + seq_len].to(torch.float64).reshape(seq_len, n, -1, 2) - ) - x_i = torch.view_as_real(x_i * freqs_i).flatten(2) - output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = x_i - seq_bucket.append(seq_bucket[-1] + seq_len) - return output.to(input_dtype) - - -@torch.compiler.disable -@torch.amp.autocast(current_omni_platform.device_type, enabled=False) -def rope_apply_s2v(x, grid_sizes, freqs, start=None): - """Apply RoPE using precomputed complex freqs (for S2V main transformer). - - Under TP, x has fewer heads than freqs (which is precomputed with all heads). - Since RoPE frequencies are identical across heads, we slice freqs to match. - """ - n = x.size(2) - input_dtype = x.dtype - output = [] - for i, _ in enumerate(x): - s = x.size(1) - x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, n, -1, 2)) - freqs_i = freqs[i, :s, :n] - x_i = torch.view_as_real(x_i * freqs_i).flatten(2) - x_i = torch.cat([x_i, x[i, s:]]) - output.append(x_i) - return torch.stack(output).to(input_dtype) - - -@torch.amp.autocast(current_omni_platform.device_type, enabled=False) -def rope_apply_usp(x, grid_sizes, freqs): - """Apply RoPE for Ulysses sequence parallel (context parallel mode).""" - n = x.size(2) - input_dtype = x.dtype - output = [] - for i, _ in enumerate(x): - s = x.size(1) - x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, n, -1, 2)) - freqs_i = freqs[i, :, :n] - x_i = torch.view_as_real(x_i * freqs_i).flatten(2) - x_i = torch.cat([x_i, x[i, s:]]) - output.append(x_i) - return torch.stack(output).to(input_dtype) - - def zero_module(module): for p in module.parameters(): p.detach().zero_() @@ -292,6 +129,8 @@ class WanS2VSelfAttention(nn.Module): causal=False, ) + self.rotary_embedding = RotaryEmbeddingWanS2V() + def forward(self, x, seq_lens, grid_sizes, freqs): b, s = x.shape[:2] @@ -311,15 +150,12 @@ class WanS2VSelfAttention(nn.Module): value = value.view(b, s, self.num_kv_heads, self.head_dim) # Apply S2V complex-valued RoPE - query = rope_apply_s2v(query, grid_sizes, freqs) - key = rope_apply_s2v(key, grid_sizes, freqs) + query = self.rotary_embedding(query, freqs) + key = self.rotary_embedding(key, freqs) - # Reshape back after rope (rope_apply_s2v preserves [B, S, N, D]) + # Reshape back after RoPE (preserves [B, S, N, D]) query = query.view(b, s, self.num_heads, self.head_dim) key = key.view(b, s, self.num_kv_heads, self.head_dim) - - # Attention — rope_apply_s2v preserves input dtype (bf16), so - # query/key/value are all bf16 here, matching FlashAttention requirements. hidden_states = self.attn(query, key, value) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(x) @@ -411,7 +247,24 @@ class WanS2VTransformerBlock(nn.Module): # 6-way scale-shift modulation self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens): + def forward( + self, + hidden_states, + encoder_hidden_states=None, + seq_lens=None, + grid_sizes=None, + freqs=None, + context=None, + context_lens=None, + e=None, + ): + x = hidden_states + if encoder_hidden_states is None: + encoder_hidden_states = e + if encoder_hidden_states is None: + raise ValueError("S2V transformer block requires modulation states.") + + e = encoder_hidden_states seg_boundary = e[1] seg_idx = [0, min(max(0, seg_boundary), x.size(1)), x.size(1)] e = e[0] @@ -450,25 +303,68 @@ class WanS2VTransformerBlock(nn.Module): return x -class WanS2VHead(nn.Module): - """S2V output head with modulation.""" +class WanS2VConditionEmbedder(nn.Module): + """Combined time and text condition embeddings for S2V.""" - def __init__(self, dim, out_dim, patch_size, eps=1e-6): + def __init__(self, dim: int, freq_dim: int, text_dim: int, text_len: int): + super().__init__() + self.text_len = text_len + self.timesteps_proj = Timesteps(num_channels=freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=freq_dim, time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, dim * 6) + self.text_embedder = PixArtAlphaTextProjection(text_dim, dim, act_fn="gelu_tanh") + + def forward( + self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, hidden_states_dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + timestep = self.timesteps_proj(timestep) + temb = self.time_embedder(timestep.to(hidden_states_dtype)) + timestep_proj = self.time_proj(self.act_fn(temb)) + encoder_hidden_states = self.text_embedder( + torch.stack( + [torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in encoder_hidden_states] + ) + ) + return temb, timestep_proj, encoder_hidden_states + + +class WanS2VTimestepProjPrepare(nn.Module): + """Prepares timestep_proj with segment modulation for S2V.""" + + def __init__(self, dim: int): super().__init__() self.dim = dim - self.out_dim = out_dim - self.patch_size = patch_size - out_dim_full = math.prod(patch_size) * out_dim - self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.head = nn.Linear(dim, out_dim_full) - self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + def forward(self, timestep_proj: torch.Tensor, temb: torch.Tensor, zero_timestep: bool, seg_boundary: int): + timestep_proj = timestep_proj.unflatten(1, (6, self.dim)) + if zero_timestep: + temb = temb[:-1] + zero_timestep_proj = timestep_proj[-1:] + timestep_proj = timestep_proj[:-1] + timestep_proj = torch.cat( + [ + timestep_proj.unsqueeze(2), + zero_timestep_proj.unsqueeze(2).repeat(timestep_proj.size(0), 1, 1, 1), + ], + dim=2, + ) + return temb, [timestep_proj, seg_boundary] + else: + timestep_proj = timestep_proj.unsqueeze(2).repeat(1, 1, 2, 1) + return temb, [timestep_proj, 0] - def forward(self, x, e): - e = (self.modulation.to(e.dtype) + e.unsqueeze(1)).chunk(2, dim=1) - x = self.norm(x).type_as(x) * (1 + e[1]) + e[0] - x = self.head(x) - return x + +class WanS2VOutputScaleShiftPrepare(nn.Module): + """Prepares output scale/shift from temb for S2V output modulation.""" + + def __init__(self, dim: int): + super().__init__() + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + shift, scale = (self.scale_shift_table.to(temb.dtype) + temb.unsqueeze(1)).chunk(2, dim=1) + return shift, scale # --------------------------------------------------------------------------- @@ -573,43 +469,10 @@ class CausalAudioEncoder(nn.Module): return res -class AudioCrossAttention(nn.Module): - """Audio cross-attention using nn.Linear (no TP needed).""" +class AudioCrossAttention(WanS2VCrossAttention): + """Audio cross-attention — same as WanS2VCrossAttention (TP-aware).""" - def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): - if dim % num_heads != 0: - raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") - super().__init__() - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - - self.q = nn.Linear(dim, dim) - self.k = nn.Linear(dim, dim) - self.v = nn.Linear(dim, dim) - self.o = nn.Linear(dim, dim) - self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - - self.attn = Attention( - num_heads=num_heads, - head_size=self.head_dim, - num_kv_heads=num_heads, - softmax_scale=1.0 / (self.head_dim**0.5), - causal=False, - ) - - def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: list[int] | None = None) -> torch.Tensor: - b, n, d = x.size(0), self.num_heads, self.head_dim - - q = self.norm_q(self.q(x)).view(b, -1, n, d) - k = self.norm_k(self.k(context)).view(b, -1, n, d) - v = self.v(context).view(b, -1, n, d) - - x = self.attn(q, k, v) - x = x.flatten(2, 3) - x = self.o(x) - return x + pass class AudioInjector_WAN(nn.Module): @@ -692,8 +555,15 @@ class SimpleSelfAttention(nn.Module): causal=False, ) + self.rotary_embedding = RotaryEmbeddingS2VGrid() + def forward( - self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor + self, + x: torch.Tensor, + seq_lens: list[int], + grid_sizes: list[tuple[int, int, int]], + freqs: torch.Tensor, + precomputed_freqs: torch.Tensor | None = None, ) -> torch.Tensor: b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim @@ -701,10 +571,13 @@ class SimpleSelfAttention(nn.Module): k = self.norm_k(self.k(x)).view(b, s, n, d) v = self.v(x).view(b, s, n, d) - q = rope_apply(q, grid_sizes, freqs) - k = rope_apply(k, grid_sizes, freqs) + if precomputed_freqs is not None: + q = RotaryEmbeddingS2VGrid.apply_precomputed(q, precomputed_freqs) + k = RotaryEmbeddingS2VGrid.apply_precomputed(k, precomputed_freqs) + else: + q = self.rotary_embedding(q, grid_sizes, freqs) + k = self.rotary_embedding(k, grid_sizes, freqs) - # rope_apply preserves [B, S, N, D] shape and input dtype q = q.view(b, s, n, d) k = k.view(b, s, n, d) @@ -716,45 +589,54 @@ class SimpleSelfAttention(nn.Module): class SwinSelfAttention(SimpleSelfAttention): def forward( - self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor + self, + x: torch.Tensor, + seq_lens: list[int], + grid_sizes: list[tuple[int, int, int]], + freqs: torch.Tensor, + precomputed_freqs: torch.Tensor | None = None, ) -> torch.Tensor: b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim - if b != 1: - raise ValueError( - f"SwinSelfAttention only supports batch_size=1, got batch_size={b}. " - "Batched inference requires refactoring the frame reference and rearrange logic." - ) q = self.norm_q(self.q(x)).view(b, s, n, d) k = self.norm_k(self.k(x)).view(b, s, n, d) v = self.v(x).view(b, s, n, d) - q = rope_apply(q, grid_sizes, freqs) - k = rope_apply(k, grid_sizes, freqs) + if precomputed_freqs is not None: + q = RotaryEmbeddingS2VGrid.apply_precomputed(q, precomputed_freqs) + k = RotaryEmbeddingS2VGrid.apply_precomputed(k, precomputed_freqs) + else: + q = self.rotary_embedding(q, grid_sizes, freqs) + k = self.rotary_embedding(k, grid_sizes, freqs) q = q.view(b, s, n, d) k = k.view(b, s, n, d) T, H, W = grid_sizes[0].tolist() - q = rearrange(q, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) - k = rearrange(k, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) - v = rearrange(v, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) + q = rearrange(q, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W) + k = rearrange(k, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W) + v = rearrange(v, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W) - # Skip ref_q (last frame query) - unused in current implementation - q = q[:-1] + q_frames = q[:, :-1] + ref_k = k[:, -1:].expand(-1, T - 1, -1, -1, -1) + ref_v = v[:, -1:].expand(-1, T - 1, -1, -1, -1) - ref_k = repeat(k[-1:], "1 s n d -> t s n d", t=k.shape[0] - 1) - k = k[:-1] - k = torch.cat([k[:1], k, k[-1:]]) - k = torch.cat([k[1:-1], k[2:], k[:-2], ref_k], dim=1) + k_frames = k[:, :-1] + k_pad = torch.cat([k_frames[:, :1], k_frames, k_frames[:, -1:]], dim=1) + k_frames = torch.cat([k_pad[:, 1:-1], k_pad[:, 2:], k_pad[:, :-2], ref_k], dim=2) - ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=v.shape[0] - 1) - v = v[:-1] - v = torch.cat([v[:1], v, v[-1:]]) - v = torch.cat([v[1:-1], v[2:], v[:-2], ref_v], dim=1) + v_frames = v[:, :-1] + v_pad = torch.cat([v_frames[:, :1], v_frames, v_frames[:, -1:]], dim=1) + v_frames = torch.cat([v_pad[:, 1:-1], v_pad[:, 2:], v_pad[:, :-2], ref_v], dim=2) - out = self.attn(q, k, v) - out = torch.cat([out, ref_v[:1]], axis=0) - out = rearrange(out, "(b t) (h w) n d -> b (t h w) n d", t=T, h=H, w=W) + q_flat = rearrange(q_frames, "b t s n d -> (b t) s n d") + k_flat = rearrange(k_frames, "b t s n d -> (b t) s n d") + v_flat = rearrange(v_frames, "b t s n d -> (b t) s n d") + + out = self.attn(q_flat, k_flat, v_flat) + out = rearrange(out, "(b t) s n d -> b t s n d", b=b) + ref_out = v[:, -1:] + out = torch.cat([out, ref_out], dim=1) + out = rearrange(out, "b t (h w) n d -> b (t h w) n d", h=H, w=W) x = out.flatten(2, 3) x = self.o(x) return x @@ -766,11 +648,6 @@ class CausalSelfAttention(SimpleSelfAttention): ) -> torch.Tensor: shifting = 3 b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim - if b != 1: - raise ValueError( - f"CausalSelfAttention only supports batch_size=1, got batch_size={b}. " - "Batched inference requires refactoring the causal masking and frame reference logic." - ) q = self.norm_q(self.q(x)).view(b, s, n, d) k = self.norm_k(self.k(x)).view(b, s, n, d) @@ -778,55 +655,61 @@ class CausalSelfAttention(SimpleSelfAttention): T, H, W = grid_sizes[0].tolist() - q = rearrange(q, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) - k = rearrange(k, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) - v = rearrange(v, "b (t h w) n d -> (b t) (h w) n d", t=T, h=H, w=W) + q = rearrange(q, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W) + k = rearrange(k, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W) + v = rearrange(v, "b (t h w) n d -> b t (h w) n d", t=T, h=H, w=W) - # Skip ref_q (last frame query) - unused in current causal implementation - q = q[:-1] + num_frames = T - 1 - grid_sizes_q = torch.tensor([[1, H, W]] * q.shape[0], dtype=torch.long) - start = [[shifting, 0, 0]] * q.shape[0] - q = rope_apply(q, grid_sizes_q, freqs, start=start) - q = q.view(q.shape[0], -1, n, d) + # Query: skip ref frame, apply RoPE per frame + q_frames = rearrange(q[:, :-1], "b t s n d -> (b t) s n d") + grid_sizes_q = torch.tensor([[1, H, W]] * (b * num_frames), dtype=torch.long) + start = [[shifting, 0, 0]] * (b * num_frames) + q_frames = self.rotary_embedding(q_frames, grid_sizes_q, freqs, start=start) + q_frames = q_frames.view(b * num_frames, -1, n, d) - ref_k = k[-1:] - grid_sizes_ref = torch.tensor([[1, H, W]], dtype=torch.long) - start_ref = [[shifting + 10, 0, 0]] - ref_k = rope_apply(ref_k, grid_sizes_ref, freqs, start_ref) - ref_k = ref_k.view(1, -1, n, d) - ref_k = repeat(ref_k, "1 s n d -> t s n d", t=k.shape[0] - 1) + # Ref key: one per batch, RoPE with shifted start, expand to all frames + ref_k = rearrange(k[:, -1:], "b 1 s n d -> b s n d") + grid_sizes_ref = torch.tensor([[1, H, W]] * b, dtype=torch.long) + start_ref = [[shifting + 10, 0, 0]] * b + ref_k = self.rotary_embedding(ref_k, grid_sizes_ref, freqs, start_ref) + ref_k = ref_k.view(b, 1, -1, n, d).expand(-1, num_frames, -1, -1, -1) + ref_k = rearrange(ref_k, "b t s n d -> (b t) s n d") - k = k[:-1] - k = torch.cat([*([k[:1]] * shifting), k]) + # Causal key: pad with shifting duplicates, gather shifting+1 neighboring frames + k_frames = k[:, :-1] + k_pad = torch.cat([k_frames[:, :1].expand(-1, shifting, -1, -1, -1), k_frames], dim=1) cat_k = [] for i in range(shifting): - cat_k.append(k[i : i - shifting]) - cat_k.append(k[shifting:]) - k = torch.cat(cat_k, dim=1) + cat_k.append(k_pad[:, i : i + num_frames]) + cat_k.append(k_pad[:, shifting:]) + k_causal = torch.cat(cat_k, dim=2) + k_causal = rearrange(k_causal, "b t s n d -> (b t) s n d") - grid_sizes_k = torch.tensor([[shifting + 1, H, W]] * (q.shape[0]), dtype=torch.long) - k = rope_apply(k, grid_sizes_k, freqs) - k = k.view(k.shape[0], -1, n, d) - k = torch.cat([k, ref_k], dim=1) + grid_sizes_k = torch.tensor([[shifting + 1, H, W]] * (b * num_frames), dtype=torch.long) + k_causal = self.rotary_embedding(k_causal, grid_sizes_k, freqs) + k_causal = k_causal.view(b * num_frames, -1, n, d) + k_causal = torch.cat([k_causal, ref_k], dim=1) - ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=q.shape[0]) - v = v[:-1] - v = torch.cat([*([v[:1]] * shifting), v]) + # Causal value: same padding structure as key, plus ref + ref_v = v[:, -1:].expand(-1, num_frames, -1, -1, -1) + ref_v_flat = rearrange(ref_v, "b t s n d -> (b t) s n d") + v_frames = v[:, :-1] + v_pad = torch.cat([v_frames[:, :1].expand(-1, shifting, -1, -1, -1), v_frames], dim=1) cat_v = [] for i in range(shifting): - cat_v.append(v[i : i - shifting]) - cat_v.append(v[shifting:]) - v = torch.cat(cat_v, dim=1) - v = torch.cat([v, ref_v], dim=1) + cat_v.append(v_pad[:, i : i + num_frames]) + cat_v.append(v_pad[:, shifting:]) + v_causal = torch.cat(cat_v, dim=2) + v_causal = rearrange(v_causal, "b t s n d -> (b t) s n d") + v_causal = torch.cat([v_causal, ref_v_flat], dim=1) - outs = [] - for i in range(q.shape[0]): - out = self.attn(q[i : i + 1], k[i : i + 1], v[i : i + 1]) - outs.append(out) - out = torch.cat(outs, dim=0) - out = torch.cat([out, ref_v[:1]], axis=0) - out = rearrange(out, "(b t) (h w) n d -> b (t h w) n d", t=T, h=H, w=W) + # Batched attention (all frames at once) + out = self.attn(q_frames, k_causal, v_causal) + out = rearrange(out, "(b t) s n d -> b t s n d", b=b) + ref_out = v[:, -1:] + out = torch.cat([out, ref_out], dim=1) + out = rearrange(out, "b t (h w) n d -> b (t h w) n d", h=H, w=W) x = out.flatten(2, 3) x = self.o(x) return x @@ -861,9 +744,14 @@ class MotionerAttentionBlock(nn.Module): self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)) def forward( - self, x: torch.Tensor, seq_lens: list[int], grid_sizes: list[tuple[int, int, int]], freqs: torch.Tensor + self, + x: torch.Tensor, + seq_lens: list[int], + grid_sizes: list[tuple[int, int, int]], + freqs: torch.Tensor, + precomputed_freqs: torch.Tensor | None = None, ) -> torch.Tensor: - y = self.self_attn(self.norm1(x).type_as(x), seq_lens, grid_sizes, freqs) + y = self.self_attn(self.norm1(x).type_as(x), seq_lens, grid_sizes, freqs, precomputed_freqs=precomputed_freqs) x = x + y y = self.ffn(self.norm2(x).type_as(x)) x = x + y @@ -934,11 +822,8 @@ class MotionerTransformers(nn.Module): raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") if (dim // num_heads) % 2 != 0: raise ValueError(f"dim // num_heads ({dim // num_heads}) must be even") - d = dim // num_heads - self.freqs = torch.cat( - [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))], - dim=1, - ) + self.rope = WanS2VRotaryPosEmbed(num_heads, dim // num_heads) + self.freqs = self.rope.freqs self.gradient_checkpointing = False self.motion_side_len = int(math.sqrt(motion_token_num)) @@ -948,6 +833,7 @@ class MotionerTransformers(nn.Module): self.trainable_token_pos_emb = trainable_token_pos_emb if trainable_token_pos_emb: + d = dim // num_heads x = torch.zeros([1, motion_token_num, num_heads, d]) x[..., ::2] = 1 gride_sizes = [ @@ -957,13 +843,13 @@ class MotionerTransformers(nn.Module): torch.tensor([1, self.motion_side_len, self.motion_side_len]).unsqueeze(0).repeat(1, 1), ] ] - token_freqs = rope_apply(x, gride_sizes, self.freqs) + token_freqs = RotaryEmbeddingS2VGrid.forward(x, gride_sizes, self.rope.freqs) token_freqs = token_freqs[0, :, 0].reshape(motion_token_num, -1, 2) token_freqs = token_freqs * 0.01 self.token_freqs = torch.nn.Parameter(token_freqs) def forward(self, x: torch.Tensor) -> torch.Tensor: - # x[0].shape[1] contains motion frames count, but not used directly here + # x: [B, C, T, H, W] tensor device = self.patch_embedding.weight.device freqs = self.freqs if freqs.device != device: @@ -977,18 +863,20 @@ class MotionerTransformers(nn.Module): if self.enable_tsm: sample_idx = [ - _sample_indices(u.shape[1], stride=self.motion_stride, expand_ratio=self.expand_ratio, c=self.sample_c) - for u in x + _sample_indices( + x[i].shape[1], stride=self.motion_stride, expand_ratio=self.expand_ratio, c=self.sample_c + ) + for i in range(x.shape[0]) ] - x = [torch.flip(torch.flip(u, [1])[:, idx], [1]) for idx, u in zip(sample_idx, x)] + x = torch.stack([torch.flip(torch.flip(x[i], [1])[:, idx], [1]) for i, idx in enumerate(sample_idx)]) - x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + x = self.patch_embedding(x) - seq_f, seq_h, seq_w = x[0].shape[-3:] - batch_size = len(x) + batch_size = x.shape[0] + seq_f, seq_h, seq_w = x.shape[-3:] if not self.enable_tsm: - grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) - grid_sizes = [[torch.zeros_like(grid_sizes), grid_sizes, grid_sizes]] + grid_sizes_t = torch.tensor([list(x.shape[2:])], dtype=torch.long).repeat(batch_size, 1) + grid_sizes = [[torch.zeros_like(grid_sizes_t), grid_sizes_t, grid_sizes_t]] seq_f = 0 else: grid_sizes = [] @@ -1003,10 +891,9 @@ class MotionerTransformers(nn.Module): grid_sizes += tsm_frame_grid_sizes seq_f = sample_idx[0][-1] + 1 - x = [u.flatten(2).transpose(1, 2) for u in x] - seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) - x = torch.cat([u for u in x]) - batch_size = len(x) + x = x.flatten(2).transpose(1, 2) + seq_lens = torch.full((batch_size,), x.shape[1], dtype=torch.long) + batch_size = x.shape[0] token_grid_sizes = [ [ @@ -1022,11 +909,15 @@ class MotionerTransformers(nn.Module): grid_sizes = grid_sizes + token_grid_sizes token_len = self.token.shape[1] - token = self.token.clone().repeat(x.shape[0], 1, 1).contiguous() - seq_lens = seq_lens + torch.tensor([t.size(0) for t in token], dtype=torch.long) + token = self.token.clone().repeat(batch_size, 1, 1).contiguous() + seq_lens = seq_lens + token_len x = torch.cat([x, token], dim=1) - kwargs = dict(seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=freqs) + precomputed_freqs = RotaryEmbeddingS2VGrid.precompute( + x.shape[1], self.num_heads, self.dim // self.num_heads, grid_sizes, freqs, x.device + ) + + kwargs = dict(seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=freqs, precomputed_freqs=precomputed_freqs) for block in self.blocks: x = block(x, **kwargs) @@ -1062,11 +953,8 @@ class FramePackMotioner(nn.Module): raise ValueError(f"inner_dim ({inner_dim}) must be divisible by num_heads ({num_heads})") if (inner_dim // num_heads) % 2 != 0: raise ValueError(f"inner_dim // num_heads ({inner_dim // num_heads}) must be even") - d = inner_dim // num_heads - self.freqs = torch.cat( - [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))], - dim=1, - ) + self.rope = WanS2VRotaryPosEmbed(num_heads, inner_dim // num_heads) + self.freqs = self.rope.freqs self.drop_mode = drop_mode def forward(self, motion_latents: torch.Tensor, add_last_motion: int = 2) -> torch.Tensor: @@ -1147,12 +1035,7 @@ class FramePackMotioner(nn.Module): grid_sizes = grid_sizes + grid_sizes_2x + grid_sizes_4x - motion_rope_emb = rope_precompute( - motion_lat.detach().view(1, motion_lat.shape[1], self.num_heads, self.inner_dim // self.num_heads), - grid_sizes, - self.freqs, - start=None, - ) + motion_rope_emb = self.rope(motion_lat.detach(), grid_sizes) mot.append(motion_lat) mot_remb.append(motion_rope_emb) @@ -1264,9 +1147,7 @@ class WanS2VTransformer3DModel(nn.Module): # Embeddings self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) - self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim)) - self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) - self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + self.condition_embedder = WanS2VConditionEmbedder(dim, freq_dim, text_dim, text_len) # Transformer blocks (TP-enabled) self.blocks = nn.ModuleList( @@ -1277,18 +1158,18 @@ class WanS2VTransformer3DModel(nn.Module): ) # Output head - self.head = WanS2VHead(dim, out_dim, patch_size, eps) + self.norm_out = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(dim, math.prod(patch_size) * out_dim) + self.output_scale_shift_prepare = WanS2VOutputScaleShiftPrepare(dim) + self.timestep_proj_prepare = WanS2VTimestepProjPrepare(dim) # RoPE base frequencies (complex-valued) if dim % num_heads != 0: raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") if (dim // num_heads) % 2 != 0: raise ValueError(f"dim // num_heads ({dim // num_heads}) must be even") - d = dim // num_heads - self.freqs = torch.cat( - [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))], - dim=1, - ) + self.rope = WanS2VRotaryPosEmbed(num_heads, dim // num_heads) + self.freqs = self.rope.freqs self.use_context_parallel = False @@ -1323,6 +1204,8 @@ class WanS2VTransformer3DModel(nn.Module): self.zero_timestep = zero_timestep self.add_last_motion = add_last_motion + self._cached_rope_key = None + self._cached_rope_emb = None # Motion modules if enable_motioner and enable_framepack: @@ -1369,7 +1252,7 @@ class WanS2VTransformer3DModel(nn.Module): .repeat(1, 1), ] ] - token_freqs = rope_apply(x, gride_sizes, self.freqs) + token_freqs = RotaryEmbeddingS2VGrid.forward(x, gride_sizes, self.rope.freqs) token_freqs = token_freqs[0, :, 0].reshape(motion_token_num, -1, 2) token_freqs = token_freqs * 0.01 self.token_freqs = nn.Parameter(token_freqs) @@ -1386,12 +1269,26 @@ class WanS2VTransformer3DModel(nn.Module): if hasattr(self, "cond_encoder"): self.cond_encoder = zero_module(self.cond_encoder) for i in range(len(self.audio_injector.injector)): - self.audio_injector.injector[i].o = zero_module(self.audio_injector.injector[i].o) + self.audio_injector.injector[i].to_out = zero_module(self.audio_injector.injector[i].to_out) if self.enable_adain: self.audio_injector.injector_adain_layers[i].linear = zero_module( self.audio_injector.injector_adain_layers[i].linear ) + # ------------------------------------------------------------------ + # RoPE preparation + # ------------------------------------------------------------------ + + def prepare_rope(self, x, grid_sizes): + """Precompute RoPE for the concatenated hidden states (cached across steps).""" + cache_key = (x.shape[0], x.shape[1], x.device) + if self._cached_rope_key == cache_key and self._cached_rope_emb is not None: + return self._cached_rope_emb + result = self.rope(x.detach(), grid_sizes) + self._cached_rope_key = cache_key + self._cached_rope_emb = result + return result + # ------------------------------------------------------------------ # Motion processing # ------------------------------------------------------------------ @@ -1414,12 +1311,7 @@ class WanS2VTransformer3DModel(nn.Module): torch.tensor([self.lat_motion_frames, height, width]).unsqueeze(0).repeat(1, 1), ] ] - motion_rope_emb = rope_precompute( - flat_mot.detach().view(1, flat_mot.shape[1], self.num_heads, self.dim // self.num_heads), - motion_grid_sizes, - self.freqs, - start=None, - ) + motion_rope_emb = self.rope(flat_mot.detach(), motion_grid_sizes) mot_remb.append(motion_rope_emb) flatten_mot.append(flat_mot) return flatten_mot, mot_remb @@ -1432,25 +1324,21 @@ class WanS2VTransformer3DModel(nn.Module): return flatten_mot, mot_remb def process_motion_transformer_motioner(self, motion_latents, drop_motion_frames=False, add_last_motion=True): - batch_size = len(motion_latents) - height = motion_latents[0].shape[2] // self.patch_size[1] - width = motion_latents[0].shape[3] // self.patch_size[2] + batch_size = motion_latents.shape[0] + height = motion_latents.shape[3] // self.patch_size[1] + width = motion_latents.shape[4] // self.patch_size[2] - freqs = self.freqs - device = self.patch_embedding.weight.device - if freqs.device != device: - freqs = freqs.to(device) + trainable_f = None if self.trainable_token_pos_emb: with torch.amp.autocast(current_omni_platform.device_type, dtype=torch.float64): token_freqs = self.token_freqs.to(torch.float64) token_freqs = token_freqs / token_freqs.norm(dim=-1, keepdim=True) - freqs = [freqs, torch.view_as_complex(token_freqs)] + trainable_f = torch.view_as_complex(token_freqs) if not drop_motion_frames and add_last_motion: - last_motion_latent = [u[:, -1:] for u in motion_latents] - last_mot = [self.patch_embedding(m.unsqueeze(0)) for m in last_motion_latent] - last_mot = [m.flatten(2).transpose(1, 2) for m in last_mot] - last_mot = torch.cat(last_mot) + last_motion_latent = motion_latents[:, :, -1:] + last_mot = self.patch_embedding(last_motion_latent) + last_mot = last_mot.flatten(2).transpose(1, 2) gride_sizes = [ [ torch.tensor([-1, 0, 0]).unsqueeze(0).repeat(batch_size, 1), @@ -1459,9 +1347,7 @@ class WanS2VTransformer3DModel(nn.Module): ] ] else: - last_mot = torch.zeros( - [batch_size, 0, self.dim], device=motion_latents[0].device, dtype=motion_latents[0].dtype - ) + last_mot = torch.zeros([batch_size, 0, self.dim], device=motion_latents.device, dtype=motion_latents.dtype) gride_sizes = [] zip_motion = self.motioner(motion_latents) @@ -1483,12 +1369,7 @@ class WanS2VTransformer3DModel(nn.Module): mot = torch.cat([last_mot, zip_motion], dim=1) gride_sizes = gride_sizes + zip_motion_grid_sizes - motion_rope_emb = rope_precompute( - mot.detach().view(batch_size, mot.shape[1], self.num_heads, self.dim // self.num_heads), - gride_sizes, - freqs, - start=None, - ) + motion_rope_emb = self.rope(mot.detach(), gride_sizes, trainable_freqs=trainable_f) return [m.unsqueeze(0) for m in mot], [r.unsqueeze(0) for r in motion_rope_emb] def inject_motion( @@ -1521,12 +1402,11 @@ class WanS2VTransformer3DModel(nn.Module): audio_emb = self.merged_audio_emb num_frames = audio_emb.shape[1] - input_hidden_states = hidden_states[:, : self.original_seq_len].clone() + input_hidden_states = hidden_states[:, : self.original_seq_len] input_hidden_states = rearrange(input_hidden_states, "b (t n) c -> (b t) n c", t=num_frames) if self.enable_adain and self.adain_mode == "attn_norm": - audio_emb_global = self.audio_emb_global - audio_emb_global = rearrange(audio_emb_global, "b t n c -> (b t) n c") + audio_emb_global = rearrange(self.audio_emb_global, "b t n c -> (b t) n c") adain_hidden_states = self.audio_injector.injector_adain_layers[audio_attn_id]( input_hidden_states, temb=audio_emb_global[:, 0] ) @@ -1534,8 +1414,7 @@ class WanS2VTransformer3DModel(nn.Module): else: attn_hidden_states = self.audio_injector.injector_pre_norm_feat[audio_attn_id](input_hidden_states) - audio_emb = rearrange(audio_emb, "b t n c -> (b t) n c", t=num_frames) - attn_audio_emb = audio_emb + attn_audio_emb = rearrange(audio_emb, "b t n c -> (b t) n c") residual_out = self.audio_injector.injector[audio_attn_id]( x=attn_hidden_states, context=attn_audio_emb, @@ -1543,7 +1422,7 @@ class WanS2VTransformer3DModel(nn.Module): * attn_audio_emb.shape[1], ) residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames) - hidden_states[:, : self.original_seq_len] = hidden_states[:, : self.original_seq_len] + residual_out + hidden_states[:, : self.original_seq_len].add_(residual_out) return hidden_states @@ -1570,7 +1449,7 @@ class WanS2VTransformer3DModel(nn.Module): result = {} if self.enable_adain: audio_emb_global, audio_emb = audio_emb_res - result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :].clone() + result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :] else: audio_emb = audio_emb_res result["audio_emb"] = audio_emb[:, motion_frames[1] :, :] @@ -1578,60 +1457,51 @@ class WanS2VTransformer3DModel(nn.Module): def forward( self, - x: torch.Tensor, - t: torch.Tensor, - context: torch.Tensor, - seq_len: int, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, ref_latents: torch.Tensor, motion_latents: torch.Tensor, cond_states: torch.Tensor, - audio_input: torch.Tensor | None = None, - motion_frames: list[int] = [17, 5], + encoder_hidden_states_audio: dict[str, torch.Tensor] | None = None, add_last_motion: int = 2, drop_motion_frames: bool = False, - audio_emb: dict[str, torch.Tensor] | None = None, - *extra_args: object, + return_dict: bool = True, **extra_kwargs: object, ) -> torch.Tensor: """Forward pass for S2V transformer. Args: - x: Noisy latents [B, C, T, H, W] - t: Timesteps [B] - context: Text embeddings [B, seq_len, dim] - seq_len: Sequence length + hidden_states: Noisy latents [B, C, T, H, W] + timestep: Timesteps [B] + encoder_hidden_states: Text embeddings [B, seq_len, dim] ref_latents: Reference image latents [B, C, 1, H, W] motion_latents: Motion context latents [B, C, T_motion, H, W] - cond_states: Pose condition latents [B, C, T, H, W] — typically zero tensor - for no pose control. Iterates over batch dimension (dim 0). - audio_input: Raw audio features (optional if audio_emb provided) - motion_frames: [num_frames, num_latent_frames] for audio alignment + cond_states: Pose condition latents [B, C, T, H, W] + encoder_hidden_states_audio: Precomputed audio embeddings from encode_audio() add_last_motion: Number of last motion frames to repeat drop_motion_frames: Whether to drop first motion frame - audio_emb: Precomputed audio embeddings from encode_audio() + return_dict: Whether to return Transformer2DModelOutput Returns: Predicted noise [B, C, T, H, W] """ + # Clear cached state from previous forward (prevents stale data from + # dummy warmup or prior clips polluting this inference) + self._cached_rope_key = None + self._cached_rope_emb = None + + x = hidden_states + t = timestep + context = encoder_hidden_states + audio_emb = encoder_hidden_states_audio add_last_motion = self.add_last_motion * add_last_motion - # Audio encoding — use precomputed embeddings if available - if audio_emb is not None: - self.merged_audio_emb = audio_emb["audio_emb"] - if self.enable_adain: - self.audio_emb_global = audio_emb["audio_emb_global"] - else: - audio_input = torch.cat( - [audio_input[..., 0:1].repeat(1, 1, 1, motion_frames[0]), audio_input], - dim=-1, - ) - audio_emb_res = self.casual_audio_encoder(audio_input) - if self.enable_adain: - audio_emb_global, audio_emb_local = audio_emb_res - self.audio_emb_global = audio_emb_global[:, motion_frames[1] :].clone() - else: - audio_emb_local = audio_emb_res - self.merged_audio_emb = audio_emb_local[:, motion_frames[1] :, :] + # Audio embeddings — must be precomputed by pipeline via encode_audio() + assert audio_emb is not None, "S2V requires precomputed audio_emb (encoder_hidden_states_audio)" + self.merged_audio_emb = audio_emb["audio_emb"] + if self.enable_adain: + self.audio_emb_global = audio_emb["audio_emb_global"] # Patch embedding + condition # NOTE: x and cond_states are 5D tensors [B, C, T, H, W]. Iterating over them @@ -1674,8 +1544,7 @@ class WanS2VTransformer3DModel(nn.Module): # Precompute RoPE x = torch.cat(x) - b, s, n, d = x.size(0), x.size(1), self.num_heads, self.dim // self.num_heads - self.pre_compute_freqs = rope_precompute(x.detach().view(b, s, n, d), grid_sizes, self.freqs, start=None) + self.pre_compute_freqs = self.prepare_rope(x, grid_sizes) x = [u.unsqueeze(0) for u in x] self.pre_compute_freqs = [u.unsqueeze(0) for u in self.pre_compute_freqs] @@ -1697,29 +1566,13 @@ class WanS2VTransformer3DModel(nn.Module): x = x + self.trainable_cond_mask(mask_input).to(x.dtype) - # Time embeddings — sinusoidal_embedding_1d uses float64 internally - # for precision, then we cast to the model dtype (bf16) since the - # downstream scale/shift modulation doesn't need float32. + # Condition embeddings (time + text) if self.zero_timestep: t = torch.cat([t, torch.zeros([1], dtype=t.dtype, device=t.device)]) - e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).to(x.dtype)) - e0 = self.time_projection(e).unflatten(1, (6, self.dim)) + e, timestep_proj, context = self.condition_embedder(t, context, x.dtype) + e, e0 = self.timestep_proj_prepare(timestep_proj, e, self.zero_timestep, self.original_seq_len) - if self.zero_timestep: - e = e[:-1] - zero_e0 = e0[-1:] - e0 = e0[:-1] - e0 = torch.cat([e0.unsqueeze(2), zero_e0.unsqueeze(2).repeat(e0.size(0), 1, 1, 1)], dim=2) - e0 = [e0, self.original_seq_len] - else: - e0 = e0.unsqueeze(2).repeat(1, 1, 2, 1) - e0 = [e0, 0] - - # Context embedding context_lens = None - context = self.text_embedding( - torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]) - ) # Transformer blocks kwargs = dict( @@ -1735,11 +1588,19 @@ class WanS2VTransformer3DModel(nn.Module): x = block(x, **kwargs) x = self.after_transformer_block(idx, x) - # Output + # Output norm, projection & unpatchify x = x[:, : self.original_seq_len] - x = self.head(x, e) + shift, scale = self.output_scale_shift_prepare(e) + x = self.norm_out(x).type_as(x) * (1 + scale) + shift + x = self.proj_out(x) x = self.unpatchify(x, original_grid_sizes) - return [u.float() for u in x] + output = torch.stack([u.float() for u in x]) + + if not return_dict: + return (output,) + from diffusers.models.modeling_outputs import Transformer2DModelOutput + + return Transformer2DModelOutput(sample=output) def unpatchify(self, x, grid_sizes): c = self.out_dim @@ -1802,26 +1663,44 @@ class WanS2VTransformer3DModel(nn.Module): name = name.replace(".ffn.0.", ".ffn.net_0.proj.") name = name.replace(".ffn.2.", ".ffn.net_2.") + # Audio injector weight remapping (q/k/v/o → to_q/to_k/to_v/to_out) + if "audio_injector.injector." in name: + name = name.replace(".q.", ".to_q.") + name = name.replace(".k.", ".to_k.") + name = name.replace(".v.", ".to_v.") + name = name.replace(".o.", ".to_out.") + + # Output head remapping (head.* → norm_out/proj_out/output_scale_shift_prepare) + if name.startswith("head.head."): + name = "proj_out." + name[len("head.head.") :] + elif name == "head.modulation": + name = "output_scale_shift_prepare.scale_shift_table" + + # Condition embedder remapping (only top-level) + # time_embedding.0.* → condition_embedder.time_embedder.linear_1.* + # time_embedding.2.* → condition_embedder.time_embedder.linear_2.* + if name.startswith("time_embedding.0."): + name = "condition_embedder.time_embedder.linear_1." + name[len("time_embedding.0.") :] + elif name.startswith("time_embedding.2."): + name = "condition_embedder.time_embedder.linear_2." + name[len("time_embedding.2.") :] + # time_projection.1.* → condition_embedder.time_proj.* + elif name.startswith("time_projection.1."): + name = "condition_embedder.time_proj." + name[len("time_projection.1.") :] + # text_embedding.0.* → condition_embedder.text_embedder.linear_1.* + # text_embedding.2.* → condition_embedder.text_embedder.linear_2.* + elif name.startswith("text_embedding.0."): + name = "condition_embedder.text_embedder.linear_1." + name[len("text_embedding.0.") :] + elif name.startswith("text_embedding.2."): + name = "condition_embedder.text_embedder.linear_2." + name[len("text_embedding.2.") :] + if name not in params_dict: logger.warning(f"Skipping weight {original_name} -> {name}") continue param = params_dict[name] - # TP-shard norm weights for main blocks - if ( - tp_size > 1 - and is_main_block - and any( - norm_name in name - for norm_name in [ - ".self_attn.norm_q.", - ".self_attn.norm_k.", - ".cross_attn.norm_q.", - ".cross_attn.norm_k.", - ] - ) - ): + # TP-shard norm weights for TP-aware attention modules + if tp_size > 1 and (".norm_q." in name or ".norm_k." in name) and loaded_weight.shape[0] != param.shape[0]: shard_size = loaded_weight.shape[0] // tp_size loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size] diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 81889607a..8edf7ec19 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -69,8 +69,7 @@ class DistributedRMSNorm(nn.Module): local_count = x.shape[-1] if tp_size > 1: - global_sum_sq = local_sum_sq.clone() - tensor_model_parallel_all_reduce(global_sum_sq) + global_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq) global_count = local_count * tp_size else: global_sum_sq = local_sum_sq diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index c8da748eb..613fb64c7 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -228,7 +228,9 @@ class StageDiffusionClient(StageClientBase): elif msg_type == "error": req_id = msg.get("request_id") rpc_id = msg.get("rpc_id") - error_msg = msg.get("error") + error_msg = msg.get("error") or "Unknown diffusion subprocess error." + status_code = msg.get("status_code") + error_type = msg.get("error_type") logger.error( "[StageDiffusionClient] stage-%s [rep-%s] subprocess error for %s: %s", self.stage_id, @@ -245,7 +247,14 @@ class StageDiffusionClient(StageClientBase): # Route request errors as error outputs so the Orchestrator # sees the request complete (instead of hanging forever). if req_id is not None: - self._output_queue.put_nowait(OmniRequestOutput.from_error(req_id, error_msg)) + self._output_queue.put_nowait( + OmniRequestOutput.from_error( + req_id, + error_msg, + status_code=status_code, + error_type=error_type, + ) + ) # Fields that are subprocess-local and cannot be serialized across # process boundaries. They are recreated in the subprocess with diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 07489d4c9..27c6f16a0 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -33,6 +33,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import ( ) from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage from vllm_omni.engine.stage_init_utils import set_death_signal +from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -400,12 +401,15 @@ class StageDiffusionProc: ) except Exception as e: logger.exception("Diffusion request %s failed: %s", request_id, e) + status_code, error_type = client_error_metadata(e) await response_socket.send( encoder.encode( { "type": "error", "request_id": request_id, "error": str(e), + "status_code": status_code, + "error_type": error_type, } ) ) @@ -483,12 +487,15 @@ class StageDiffusionProc: ) except Exception as e: logger.exception("Batch diffusion request %s failed: %s", rid, e) + status_code, error_type = client_error_metadata(e) await response_socket.send( encoder.encode( { "type": "error", "request_id": rid, "error": str(e), + "status_code": status_code, + "error_type": error_type, } ) ) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 636ac7f60..55e3dca1f 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -101,8 +101,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): def load_model( self, memory_pool_context_fn: callable | None = None, - load_format: str | None = None, - custom_pipeline_name: str | type | None = None, + load_format: str = "default", + custom_pipeline_name: str | None = None, ) -> None: """ Load the diffusion model, apply compilation and offloading. @@ -138,7 +138,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): with get_memory_context(): with DeviceMemoryProfiler() as m: self.pipeline = model_loader.load_model( - od_config=self.od_config, load_device=load_device, load_format=load_format, custom_pipeline_name=custom_pipeline_name, diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 3ade1f64d..cd9576437 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -801,7 +801,7 @@ class WorkerProc: except Exception as e: logger.error(f"Error processing RPC: {e}", exc_info=True) if self.result_mq is not None: - self.return_result(DiffusionOutput(error=str(e))) + self.return_result(DiffusionOutput.from_exception(e)) elif isinstance(msg, dict) and msg.get("type") == "shutdown": logger.info("Worker %s: Received shutdown message", self.gpu_id) @@ -817,7 +817,7 @@ class WorkerProc: f"Error executing forward in event loop: {e}", exc_info=True, ) - output = DiffusionOutput(error=str(e)) + output = DiffusionOutput.from_exception(e) try: self.return_result(output) diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py index 7af13a322..3506f02f6 100644 --- a/vllm_omni/distributed/__init__.py +++ b/vllm_omni/distributed/__init__.py @@ -6,6 +6,7 @@ from .omni_connectors import ( MooncakeConnector, MooncakeStoreConnector, MooncakeTransferEngineConnector, + MoriTransferEngineConnector, OmniConnectorBase, OmniConnectorFactory, OmniTransferConfig, @@ -25,6 +26,7 @@ __all__ = [ "MooncakeConnector", # compat alias "MooncakeStoreConnector", "MooncakeTransferEngineConnector", + "MoriTransferEngineConnector", "SharedMemoryConnector", "YuanrongConnector", "YuanrongTransferEngineConnector", diff --git a/vllm_omni/distributed/omni_connectors/__init__.py b/vllm_omni/distributed/omni_connectors/__init__.py index 98e3a5c36..0cee596bb 100644 --- a/vllm_omni/distributed/omni_connectors/__init__.py +++ b/vllm_omni/distributed/omni_connectors/__init__.py @@ -17,6 +17,11 @@ try: from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector except ImportError: MooncakeTransferEngineConnector = None # RDMA deps (msgspec/zmq/mooncake) not installed + +try: + from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector +except ImportError: + MoriTransferEngineConnector = None # RDMA deps (msgspec/zmq/mori) not installed from .factory import OmniConnectorFactory from .utils.config import ConnectorSpec, OmniTransferConfig from .utils.initialization import ( @@ -44,6 +49,7 @@ __all__ = [ "MooncakeConnector", # compat alias → MooncakeStoreConnector "MooncakeStoreConnector", "MooncakeTransferEngineConnector", + "MoriTransferEngineConnector", "SharedMemoryConnector", "YuanrongConnector", "YuanrongTransferEngineConnector", diff --git a/vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py new file mode 100644 index 000000000..360f4527a --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py @@ -0,0 +1,962 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""OmniConnector backed by the Mori RDMA transfer engine. + +Implements the ``OmniConnector`` interface using Mori's ``IOEngine`` / +``MemoryDesc`` API to perform zero-copy RDMA transfers between +disaggregated prefill and decode workers. + +Notable design points: + * Remote peers must be registered via ``register_remote_engine()`` + with the peer's ``EngineDesc`` before any data transfer can proceed. + * Memory regions are described by ``MemoryDesc`` objects (serialisable + via pack/unpack) rather than raw virtual addresses. + * Transfers are dispatched through ``IOEngine.batch_write()`` and + tracked asynchronously via ``TransferStatus`` objects. + +ZMQ handshake protocol: + * **Pull request** – receiver sends ``MoriPullRequest`` (msgspec-encoded) + containing its ``EngineDesc`` and pool ``MemoryDesc`` so the sender can + register the remote engine and RDMA-write directly into the receiver's + pool at the specified offset. + * **Query request** – receiver sends ``QUERY_INFO`` prefix + ``QueryRequest`` + to check whether data is ready on the sender side (metadata-less get path). +""" + +import os +import queue +import socket +import threading +import time as _time_mod +import uuid +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import msgspec +import torch +import zmq + +from ..utils.logging import get_connector_logger +from ..utils.memory_pool import BufferAllocator, ManagedBuffer +from ..utils.serialization import OmniSerializer +from .base import OmniConnectorBase + +logger = get_connector_logger(__name__) + +try: + from mori.cpp import TransferStatus # noqa: F401 + from mori.io import ( + BackendType, + EngineDesc, + IOEngine, + IOEngineConfig, + MemoryDesc, + PollCqMode, + RdmaBackendConfig, + XgmiBackendConfig, + ) +except ImportError: + IOEngine = None + +# Supported backend types for Mori. Kept as string constants so configuration +# (YAML / CLI / dict) stays transport-agnostic. ``rdma`` uses NIC-based RDMA +# (RoCE / IB, GDR-capable); ``xgmi`` uses AMD Infinity Fabric GPU-to-GPU +# direct links and therefore requires a CUDA pool. +_SUPPORTED_BACKENDS = ("rdma", "xgmi") + +_BUFFER_TTL_SECONDS = 300 + +TRANS_DONE = b"trans_done" +TRANS_ERROR = b"trans_error" +QUERY_INFO = b"query_info" +INFO_NOT_FOUND = b"info_not_found" + + +# --------------------------------------------------------------------------- +# ZMQ message types +# --------------------------------------------------------------------------- + + +class MoriPullRequest(msgspec.Struct): + """Receiver → sender: request an RDMA write into the receiver's pool.""" + + request_id: str + engine_desc_packed: bytes + mem_desc_packed: bytes + dst_offset: int + length: int + + +class QueryRequest(msgspec.Struct): + """Receiver → sender: query whether data for *request_id* is ready.""" + + request_id: str + + +class QueryResponse(msgspec.Struct): + """Sender → receiver: metadata about a ready buffer.""" + + request_id: str + data_size: int + is_fast_path: bool + + +# --------------------------------------------------------------------------- +# Connector +# --------------------------------------------------------------------------- + + +class MoriTransferEngineConnector(OmniConnectorBase): + """OmniConnector backed by the Mori RDMA IOEngine. + + Topology: 1 sender ↔ 1 receiver per key (same as + ``MooncakeTransferEngineConnector``). + """ + + supports_raw_data: bool = True + + # ------------------------------------------------------------------ init + def __init__(self, config: dict[str, Any]): + if IOEngine is None: + raise ImportError("Mori is not available. Install via: pip install mori") + + self._closed = False + self._bind_error: Exception | None = None + + # Teardown-safe defaults + self._stop_event = threading.Event() + self._sender_executor: ThreadPoolExecutor | None = None + self._listener_thread: threading.Thread | None = None + self._listener_ready = threading.Event() + self._local_buffers: dict[str, Any] = {} + self._local_buffers_lock = threading.Lock() + self._req_local = threading.local() + self._worker_local = threading.local() + self._last_ttl_check: float = _time_mod.monotonic() + + # Track remote engines already registered with IOEngine. + # ``mori.io.IOEngine.deregister_remote_engine`` requires the original + # ``EngineDesc`` object (its pybind signature is typed + # ``deregister_remote_engine(self, engine_desc: EngineDesc)``), not + # just the key string -- so we keep the desc alongside the key so + # ``close()`` can pass a valid object back at teardown. + self._registered_engines: dict[str, EngineDesc] = {} + self._registered_engines_lock = threading.Lock() + + self._metrics = { + "puts": 0, + "gets": 0, + "bytes_transferred": 0, + "errors": 0, + "timeouts": 0, + } + + self.config = config + + # Stage id is read by OmniChunkTransferAdapter (process_pending_chunks / + # _send_single_request / _poll_single_request / ...) to decide whether + # this connector sits at the sender or receiver end of a stage pair. + # Kept parallel to ``SharedMemoryConnector`` and other OmniConnector + # implementations. Populated by ``build_stage_connectors`` via the + # ``stage_id`` key that ``get_connectors_config_for_stage`` injects. + self.stage_id = config.get("stage_id", -1) + + # ---- Host / ZMQ ---- + host_cfg = config.get("host", "127.0.0.1") + if host_cfg.lower() == "auto": + self.host = self._get_local_ip() + logger.info(f"Auto-detected local IP: {self.host}") + else: + self.host = host_cfg + self.zmq_port = config.get("zmq_port", 50051) + + # ---- Backend selection ---- + # Defaults to "rdma" so existing deployments (pre-XGMI support) keep + # working unchanged. ``xgmi`` selects AMD Infinity Fabric GPU-to-GPU + # direct links; see ``mori.io.XgmiBackendConfig``. + backend_type_cfg = str(config.get("backend_type", "rdma")).lower() + if backend_type_cfg not in _SUPPORTED_BACKENDS: + raise ValueError(f"Invalid backend_type={backend_type_cfg!r}. Supported: {list(_SUPPORTED_BACKENDS)}.") + self.backend_type = backend_type_cfg + + # ---- RDMA device (RDMA backend only) ---- + self.device_name = "" + if self.backend_type == "rdma": + self.device_name = config.get("device_name", "") + if not self.device_name: + env_dev = os.environ.get("MORI_RDMA_DEVICES", "") + if env_dev: + self.device_name = env_dev + logger.info(f"Using MORI_RDMA_DEVICES from env: {self.device_name}") + elif config.get("device_name"): + logger.warning( + "device_name=%r is ignored for backend_type='xgmi' (XGMI does not use an RDMA NIC).", + config.get("device_name"), + ) + + # ---- Pool config ---- + self.pool_size = config.get("memory_pool_size", 1024**3) + self.pool_device = config.get("memory_pool_device", "cpu") + if self.backend_type == "xgmi" and self.pool_device == "cpu": + raise ValueError( + "backend_type='xgmi' requires memory_pool_device='cuda': " + "XGMI is a GPU-to-GPU fabric and cannot address CPU memory." + ) + + # ---- Sender info (receiver uses this when metadata=None) ---- + self.sender_host = config.get("sender_host", None) + self.sender_zmq_port = config.get("sender_zmq_port", None) + + # ---- Role ---- + # ``sender`` : binds the ZMQ listener, accepts ``put()``. + # ``receiver`` : no listener, only ``get()`` by querying an + # upstream sender whose endpoint lives in + # ``sender_host`` / ``sender_zmq_port``. + role = str(config.get("role", "sender")).lower() + if role not in {"sender", "receiver"}: + raise ValueError(f"Invalid role={role!r} for MoriTransferEngineConnector. Expected 'sender' or 'receiver'.") + self.role = role + self.can_put = role == "sender" + + # ---- Mori IOEngine ---- + if self.device_name: + os.environ["MORI_RDMA_DEVICES"] = self.device_name + + engine_config = IOEngineConfig(host=self.host, port=0) + self.engine_key = f"omni-{role}-{uuid.uuid4().hex[:8]}-pid{os.getpid()}-{self.host}" + self.engine = IOEngine(self.engine_key, engine_config) + + # ---- Backend creation (per backend_type) ---- + # ``IOEngine.batch_write`` is backend-agnostic, so only construction + # differs; the data-plane / ZMQ handshake paths below are identical. + if self.backend_type == "xgmi": + xgmi_cfg = XgmiBackendConfig() + xgmi_cfg.num_streams = config.get("xgmi_num_streams", 64) + xgmi_cfg.num_events = config.get("xgmi_num_events", 64) + self.engine.create_backend(BackendType.XGMI, xgmi_cfg) + logger.info(f"Mori backend: XGMI (num_streams={xgmi_cfg.num_streams}, num_events={xgmi_cfg.num_events})") + else: # rdma + qp_per_transfer = config.get("qp_per_transfer", 1) + post_batch_size = config.get("post_batch_size", -1) + num_workers = config.get("num_worker_threads", 1) + rdma_cfg = RdmaBackendConfig( + qp_per_transfer, + post_batch_size, + num_workers, + PollCqMode.POLLING, + False, + ) + self.engine.create_backend(BackendType.RDMA, rdma_cfg) + logger.info( + f"Mori backend: RDMA (qp_per_transfer={qp_per_transfer}, " + f"num_workers={num_workers}, device={self.device_name or 'auto'})" + ) + + self.engine_desc: EngineDesc = self.engine.get_engine_desc() + self.engine_desc_packed: bytes = self.engine_desc.pack() + logger.info(f"Mori IOEngine ready: key={self.engine_key} at {self.engine_desc.host}:{self.engine_desc.port}") + + # ---- Pool allocation & Mori memory registration ---- + logger.info( + f"Allocating {self.backend_type.upper()} pool: {self.pool_size / 1024**2:.2f} MB on {self.pool_device}" + ) + try: + if self.pool_device == "cpu": + self.pool = torch.empty(self.pool_size, dtype=torch.uint8).pin_memory() + else: + self.pool = torch.empty( + self.pool_size, + dtype=torch.uint8, + device=self.pool_device, + ) + self.pool_mem_desc: MemoryDesc = self.engine.register_torch_tensor(self.pool) + self.pool_mem_desc_packed: bytes = self.pool_mem_desc.pack() + self.base_ptr = self.pool.data_ptr() + except Exception as e: + logger.error(f"Failed to allocate/register pool: {e}") + raise + + self.allocator = BufferAllocator(self.pool_size, alignment=4096) + + # ---- ZMQ / threading ---- + self.zmq_ctx = zmq.Context() + self._sender_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="mori-sender") + + logger.info( + f"MoriTransferEngineConnector config:\n" + f" Local: host={self.host}, zmq_port={self.zmq_port}\n" + f" Remote: sender_host={self.sender_host}, " + f"sender_zmq_port={self.sender_zmq_port}\n" + f" Role: can_put={self.can_put}, " + f"configured_role={config.get('role', 'sender')}" + ) + + if self.can_put: + self._last_ttl_check = _time_mod.monotonic() + self._listener_thread = threading.Thread(target=self._zmq_listener_loop, daemon=True) + self._listener_thread.start() + self._listener_ready.wait(timeout=1.0) + if self._bind_error is not None: + raise RuntimeError( + f"MoriTransferEngineConnector failed to bind ZMQ on {self.host}:{self.zmq_port}: {self._bind_error}" + ) from self._bind_error + logger.info(f"MoriTransferEngineConnector SENDER ready (ZMQ on {self.host}:{self.zmq_port})") + else: + if not self.sender_host or self.sender_host.lower() == "auto": + logger.info("MoriTransferEngineConnector RECEIVER: awaiting sender info via update_sender_info().") + else: + logger.info( + f"MoriTransferEngineConnector RECEIVER ready (sender at {self.sender_host}:{self.sender_zmq_port})" + ) + + # -------------------------------------------------------- public helpers + def get_connection_info(self) -> dict[str, Any]: + return { + "host": self.host, + "zmq_port": self.zmq_port, + "engine_key": self.engine_key, + "can_put": self.can_put, + } + + def update_sender_info(self, sender_host: str, sender_zmq_port: int) -> None: + """Inject the sender's ZMQ endpoint into the receiver connector.""" + self.sender_host = sender_host + self.sender_zmq_port = sender_zmq_port + logger.info(f"Sender info updated: host={sender_host!r}, zmq_port={sender_zmq_port}") + + # -------------------------------------------------- internal helpers + @staticmethod + def _get_local_ip() -> str: + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: + s.connect(("8.8.8.8", 80)) + return s.getsockname()[0] + except Exception: + try: + return socket.gethostbyname(socket.gethostname()) + except Exception: + return "127.0.0.1" + + def _get_req_socket(self, zmq_addr: str, timeout_ms: int) -> zmq.Socket: + cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None) + if cache is None: + cache = {} + self._req_local.cache = cache + + sock = cache.get(zmq_addr) + if sock is None: + sock = self.zmq_ctx.socket(zmq.REQ) + sock.connect(zmq_addr) + cache[zmq_addr] = sock + sock.setsockopt(zmq.SNDTIMEO, timeout_ms) + sock.setsockopt(zmq.RCVTIMEO, timeout_ms) + return sock + + def _invalidate_req_socket(self, zmq_addr: str) -> None: + cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None) + if cache is None: + return + sock = cache.pop(zmq_addr, None) + if sock is not None: + try: + sock.close(linger=0) + except Exception: + pass + + def _ensure_remote_registered(self, engine_desc_packed: bytes) -> str: + """Register a remote engine with the local IOEngine (idempotent).""" + desc = EngineDesc.unpack(engine_desc_packed) + key = desc.key + with self._registered_engines_lock: + if key not in self._registered_engines: + self.engine.register_remote_engine(desc) + self._registered_engines[key] = desc + logger.debug(f"Registered remote Mori engine: {key}") + return key + + # ----------------------------------------------------------------- put() + def put( + self, + from_stage: str, + to_stage: str, + put_key: str, + data: Any, + ) -> tuple[bool, int, dict[str, Any] | None]: + if self._closed: + raise RuntimeError("Cannot put: MoriTransferEngineConnector is closed") + if not self.can_put: + logger.warning(f"Rejecting put for {put_key}: connector is receiver-only") + return False, 0, None + + put_key = self._make_key(put_key, from_stage, to_stage) + try: + src_offset = 0 + size = 0 + holder = None + is_fast_path = True + should_release = False + + # Reject empty payloads early + if isinstance(data, bytes) and len(data) == 0: + return False, 0, None + if isinstance(data, torch.Tensor) and data.nbytes == 0: + return False, 0, None + + # Serialize non-raw types + if not isinstance(data, (ManagedBuffer, torch.Tensor, bytes)): + data = OmniSerializer.serialize(data) + is_fast_path = False + + if isinstance(data, ManagedBuffer): + if data.pool_tensor.data_ptr() != self.pool.data_ptr(): + data = data.tensor.contiguous() + else: + src_offset = data.offset + size = data.size + holder = data + should_release = False + + if isinstance(data, (torch.Tensor, bytes)): + if isinstance(data, torch.Tensor): + size = data.nbytes + tensor_data = data + else: + size = len(data) + tensor_data = torch.frombuffer(data, dtype=torch.uint8) + + try: + offset = self.allocator.alloc(size) + holder = ManagedBuffer(self.allocator, offset, size, self.pool) + should_release = True + except MemoryError: + logger.error(f"Pool exhausted for {size} bytes") + return False, 0, None + + try: + dst_t = holder.tensor + if isinstance(data, torch.Tensor): + if not data.is_contiguous(): + data = data.contiguous() + src_view = data.view(torch.uint8).flatten() + if src_view.device != dst_t.device: + dst_t.copy_(src_view, non_blocking=True) + if src_view.is_cuda: + with torch.cuda.device(src_view.device): + torch.cuda.current_stream().synchronize() + elif dst_t.is_cuda: + with torch.cuda.device(dst_t.device): + torch.cuda.current_stream().synchronize() + else: + dst_t.copy_(src_view) + if dst_t.is_cuda: + with torch.cuda.device(dst_t.device): + torch.cuda.current_stream().synchronize() + else: + dst_t.copy_(tensor_data) + if dst_t.is_cuda: + with torch.cuda.device(dst_t.device): + torch.cuda.current_stream().synchronize() + except Exception as e: + holder.release() + logger.error(f"Failed to copy data to pool: {e}") + return False, 0, None + + src_offset = offset + + if size <= 0: + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + return False, 0, None + + with self._local_buffers_lock: + old = self._local_buffers.pop(put_key, None) + if old: + _, _, oh, osr, _, _ = old + if osr and isinstance(oh, ManagedBuffer): + oh.release() + logger.warning(f"Released stale buffer for duplicate key: {put_key}") + # (offset, size, holder, should_release, is_fast_path, ts) + self._local_buffers[put_key] = ( + src_offset, + size, + holder, + should_release, + is_fast_path, + _time_mod.monotonic(), + ) + + metadata = { + "source_host": self.host, + "source_port": self.zmq_port, + "data_size": size, + "is_fast_path": is_fast_path, + } + self._metrics["puts"] += 1 + self._metrics["bytes_transferred"] += size + return True, size, metadata + + except Exception as e: + self._metrics["errors"] += 1 + logger.error(f"Put failed for {put_key}: {e}", exc_info=True) + return False, 0, None + + # ------------------------------------------------ query (no-metadata get) + def _query_metadata_from_sender(self, get_key: str) -> dict[str, Any] | None: + zmq_addr = f"tcp://{self.sender_host}:{self.sender_zmq_port}" + req_socket = self._get_req_socket(zmq_addr, timeout_ms=5000) + try: + q = QueryRequest(request_id=get_key) + req_socket.send(QUERY_INFO + msgspec.msgpack.encode(q)) + resp = req_socket.recv() + if resp == INFO_NOT_FOUND: + return None + qr = msgspec.msgpack.decode(resp, type=QueryResponse) + return { + "source_host": self.sender_host, + "source_port": self.sender_zmq_port, + "data_size": qr.data_size, + "is_fast_path": qr.is_fast_path, + } + except Exception as e: + self._invalidate_req_socket(zmq_addr) + logger.debug(f"Metadata query failed for {get_key}: {e}") + return None + + # ----------------------------------------------------------------- get() + def get( + self, + from_stage: str, + to_stage: str, + get_key: str, + metadata: dict[str, Any] | None = None, + ) -> tuple[Any, int] | None: + if self._closed: + raise RuntimeError("Cannot get: MoriTransferEngineConnector is closed") + + get_key = self._make_key(get_key, from_stage, to_stage) + _t0 = _time_mod.perf_counter() + + # Resolve metadata + if not metadata: + if not self.sender_host or not self.sender_zmq_port or str(self.sender_host).lower() == "auto": + raise RuntimeError("get(metadata=None) requires sender info. Call update_sender_info() first.") + metadata = self._query_metadata_from_sender(get_key) + if not metadata: + return None + + _t1 = _time_mod.perf_counter() + _query_ms = (_t1 - _t0) * 1000 + + src_host = metadata.get("source_host") + src_port = metadata.get("source_port") + data_size = metadata.get("data_size", 0) + is_fast_path = metadata.get("is_fast_path", False) + + if not src_host or not src_port or str(src_host).lower() == "auto": + logger.error(f"Invalid metadata for {get_key}") + return None + if data_size == 0: + logger.warning(f"Skipping get for {get_key}: data_size is 0") + return None + + # Allocate destination buffer + try: + offset = self.allocator.alloc(data_size) + recv_buf = ManagedBuffer(self.allocator, offset, data_size, self.pool) + except MemoryError: + logger.error(f"Failed to allocate {data_size} bytes for get") + return None + + _t2 = _time_mod.perf_counter() + _alloc_ms = (_t2 - _t1) * 1000 + + # Build pull request with Mori-specific EngineDesc + MemoryDesc + pull_req = MoriPullRequest( + request_id=get_key, + engine_desc_packed=self.engine_desc_packed, + mem_desc_packed=self.pool_mem_desc_packed, + dst_offset=offset, + length=data_size, + ) + + _base_timeout_ms = 30000 + _size_timeout_ms = max(0, (data_size // (100 * 1024 * 1024))) * 5000 + _total_timeout_ms = _base_timeout_ms + _size_timeout_ms + zmq_addr = f"tcp://{src_host}:{src_port}" + req_socket = self._get_req_socket(zmq_addr, timeout_ms=_total_timeout_ms) + + try: + req_socket.send(msgspec.msgpack.encode(pull_req)) + resp = req_socket.recv() + + _t3 = _time_mod.perf_counter() + _rdma_ms = (_t3 - _t2) * 1000 + + if resp == TRANS_DONE: + if self.pool.is_cuda: + with torch.cuda.device(self.pool.device): + torch.cuda.current_stream().synchronize() + + _t4 = _time_mod.perf_counter() + _sync_ms = (_t4 - _t3) * 1000 + + if is_fast_path: + _total_ms = (_time_mod.perf_counter() - _t0) * 1000 + _mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0 + logger.info( + f"[MORI GET] {get_key}: query={_query_ms:.1f}ms, " + f"alloc={_alloc_ms:.1f}ms, rdma={_rdma_ms:.1f}ms, " + f"sync={_sync_ms:.1f}ms, total={_total_ms:.1f}ms, " + f"{_mbps:.1f} MB/s (fast_path)" + ) + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += data_size + return recv_buf, data_size + else: + try: + raw = recv_buf.to_bytes() + val = OmniSerializer.deserialize(raw) + _total_ms = (_time_mod.perf_counter() - _t0) * 1000 + _mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0 + logger.info(f"[MORI GET] {get_key}: total={_total_ms:.1f}ms, {_mbps:.1f} MB/s (deserialized)") + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += data_size + return val, data_size + finally: + recv_buf.release() + else: + self._metrics["errors"] += 1 + logger.error(f"MORI get failed: received {resp} instead of TRANS_DONE") + recv_buf.release() + return None + except Exception as e: + self._invalidate_req_socket(zmq_addr) + self._metrics["timeouts"] += 1 + logger.error(f"MORI get error: {e}", exc_info=True) + recv_buf.release() + return None + + # ------------------------------------------------------------- cleanup() + def cleanup( + self, + request_id: str, + from_stage: str | None = None, + to_stage: str | None = None, + ) -> None: + if (from_stage is None) != (to_stage is None): + raise ValueError("cleanup() requires both from_stage and to_stage, or neither.") + if from_stage is not None and to_stage is not None: + request_id = self._make_key(request_id, from_stage, to_stage) + with self._local_buffers_lock: + item = self._local_buffers.pop(request_id, None) + if item: + _, _, holder, sr, _, _ = item + if sr and isinstance(holder, ManagedBuffer): + holder.release() + + # -------------------------------------------------------------- health() + def health(self) -> dict[str, Any]: + if self._closed: + return {"status": "unhealthy", "error": "Connector is closed"} + return { + "status": "healthy", + "host": self.host, + "pool_device": self.pool_device, + "pool_size": self.pool_size, + "engine_key": self.engine_key, + **self._metrics, + } + + # --------------------------------------------------------------- close() + def close(self) -> None: + if getattr(self, "_closed", True): + return + self._closed = True + logger.info("Closing MoriTransferEngineConnector...") + + self._stop_event.set() + + if self._listener_thread is not None and self._listener_thread.is_alive(): + self._listener_thread.join(timeout=2.0) + if self._listener_thread.is_alive(): + logger.warning("Listener thread did not stop gracefully") + + if self._sender_executor is not None: + self._sender_executor.shutdown(wait=True, cancel_futures=False) + + with self._local_buffers_lock: + for _k, item in list(self._local_buffers.items()): + _, _, holder, sr, _, _ = item + if sr and isinstance(holder, ManagedBuffer): + holder.release() + self._local_buffers.clear() + + cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None) + if cache: + for _addr, sock in cache.items(): + try: + sock.close(linger=0) + except Exception: + pass + cache.clear() + + try: + if hasattr(self, "zmq_ctx"): + self.zmq_ctx.term() + except Exception as e: + logger.warning(f"Failed to terminate ZMQ context: {e}") + + # ---- Release Mori IOEngine resources ---- + # ``mori.io.IOEngine`` does not expose an explicit ``shutdown()`` / + # ``close()`` entry-point, so we best-effort deregister every resource + # we registered with it and then drop the only reference we hold so + # Python GC triggers the C++ destructor. + engine = getattr(self, "engine", None) + if engine is not None: + pool_desc = getattr(self, "pool_mem_desc", None) + if pool_desc is not None: + try: + engine.deregister_memory(pool_desc) + except Exception as e: + logger.warning(f"Mori deregister_memory failed: {e}") + self.pool_mem_desc = None # type: ignore[assignment] + + with self._registered_engines_lock: + remote_descs = list(self._registered_engines.values()) + self._registered_engines.clear() + for desc in remote_descs: + try: + engine.deregister_remote_engine(desc) + except Exception as e: + logger.debug(f"Mori deregister_remote_engine({desc.key!r}) failed: {e}") + + # Drop the only reference to IOEngine; the C++ destructor runs + # under Python GC and unwinds backend / RDMA fabric state. + self.engine = None # type: ignore[assignment] + + self.pool = None # type: ignore[assignment] + logger.info("MoriTransferEngineConnector closed.") + + # ============================================================== LISTENER + def _cleanup_stale_buffers(self) -> None: + """Reclaim buffers older than ``_BUFFER_TTL_SECONDS``. + + Prevents permanent memory leaks when a receiver crashes or times out + without ever pulling the data. + + TODO(zejwang): In extreme rare case, long transfer time, there might + exist TTL cleanup vs in-flight RDMA transfer conflict, which will be + handled in a follow-up PR. Same race is acknowledged in + ``MooncakeTransferEngineConnector._cleanup_stale_buffers``. + """ + now = _time_mod.monotonic() + with self._local_buffers_lock: + stale = [k for k, v in self._local_buffers.items() if now - v[5] > _BUFFER_TTL_SECONDS] + for k in stale: + item = self._local_buffers.pop(k) + _, _, holder, sr, _, _ = item + if sr and isinstance(holder, ManagedBuffer): + holder.release() + logger.warning(f"TTL expired ({_BUFFER_TTL_SECONDS}s): reclaimed buffer for {k}") + + def _zmq_listener_loop(self) -> None: + sock = self.zmq_ctx.socket(zmq.ROUTER) + + try: + sock.bind(f"tcp://{self.host}:{self.zmq_port}") + except zmq.ZMQError as exc: + # Any bind failure (EADDRINUSE, EADDRNOTAVAIL, EACCES, etc.) is + # fatal for a sender — fail fast so __init__ propagates the error. + # There is no silent receiver fallback; roles are explicitly + # assigned (matches MooncakeTransferEngineConnector). + logger.error(f"ZMQ bind failed on {self.host}:{self.zmq_port}: {exc} (errno={exc.errno})") + self.can_put = False + self._bind_error = exc + self._listener_ready.set() + return + + self._listener_ready.set() + + notify_addr = f"inproc://mori-notify-{id(self)}" + notify_recv = self.zmq_ctx.socket(zmq.PULL) + notify_recv.bind(notify_addr) + + poller = zmq.Poller() + poller.register(sock, zmq.POLLIN) + poller.register(notify_recv, zmq.POLLIN) + + response_queue: queue.Queue = queue.Queue() + + try: + while not self._stop_event.is_set(): + try: + events = dict(poller.poll(1000)) + + if notify_recv in events: + while True: + try: + notify_recv.recv(zmq.NOBLOCK) + except zmq.Again: + break + + while True: + try: + identity, response = response_queue.get_nowait() + sock.send_multipart([identity, b"", response]) + except queue.Empty: + break + + now = _time_mod.monotonic() + if now - self._last_ttl_check >= 10.0: + self._last_ttl_check = now + self._cleanup_stale_buffers() + + if sock in events: + frames = sock.recv_multipart() + if len(frames) >= 2: + self._sender_executor.submit( + self._handle_pull_request, + response_queue, + notify_addr, + frames[0], + frames[-1], + ) + except zmq.ContextTerminated: + break + except Exception: + logger.debug("Listener loop error", exc_info=True) + finally: + try: + notify_recv.close(linger=0) + sock.close(linger=0) + except Exception: + pass + + def _handle_pull_request( + self, + response_queue: queue.Queue, + notify_addr: str, + identity: bytes, + payload: bytes, + ) -> None: + try: + if payload.startswith(QUERY_INFO): + self._handle_query_request( + response_queue, + notify_addr, + identity, + payload[len(QUERY_INFO) :], + ) + return + + pull = msgspec.msgpack.decode(payload, type=MoriPullRequest) + + with self._local_buffers_lock: + item = self._local_buffers.get(pull.request_id) + + if not item: + response_queue.put((identity, TRANS_ERROR)) + self._notify_listener(notify_addr) + return + + src_offset, src_size, _, _, _, _ = item + + # Validate against sender's own src_size before RDMA write: + # length < src_size silently truncates the payload; length > + # src_size reads past this allocation into adjacent in-flight + # buffers because pool_mem_desc covers the whole pool, not just + # this slot, so Mori does not bound-check. Refuse on mismatch + # rather than corrupt or leak data silently. + if pull.length != src_size: + logger.error( + "Length mismatch for %s: sender src_size=%d != receiver length=%d; refusing transfer", + pull.request_id, + src_size, + pull.length, + ) + response_queue.put((identity, TRANS_ERROR)) + self._notify_listener(notify_addr) + return + + # Register the receiver's IOEngine (idempotent) + self._ensure_remote_registered(pull.engine_desc_packed) + + # Reconstruct the receiver's pool MemoryDesc + remote_mem = MemoryDesc.unpack(pull.mem_desc_packed) + + # RDMA write from local pool → remote pool. Use src_size (sender's + # own record); it has just been validated to equal pull.length. + transfer_uid = self.engine.allocate_transfer_uid() + statuses = self.engine.batch_write( + [self.pool_mem_desc], + [[src_offset]], + [remote_mem], + [[pull.dst_offset]], + [[src_size]], + [transfer_uid], + ) + + success = True + for st in statuses: + st.Wait() + if st.Failed(): + logger.error(f"RDMA write failed for {pull.request_id}: {st.Message()}") + success = False + + if success: + self.cleanup(pull.request_id) + response_queue.put((identity, TRANS_DONE)) + else: + logger.warning(f"RDMA write failed for {pull.request_id}. Buffer retained for retry.") + response_queue.put((identity, TRANS_ERROR)) + + except Exception as e: + logger.error(f"Pull request handler error: {e}", exc_info=True) + response_queue.put((identity, TRANS_ERROR)) + + self._notify_listener(notify_addr) + + def _handle_query_request( + self, + response_queue: queue.Queue, + notify_addr: str, + identity: bytes, + payload: bytes, + ) -> None: + try: + q = msgspec.msgpack.decode(payload, type=QueryRequest) + with self._local_buffers_lock: + item = self._local_buffers.get(q.request_id) + if not item: + response_queue.put((identity, INFO_NOT_FOUND)) + else: + _, data_size, _, _, is_fast, _ = item + resp = QueryResponse( + request_id=q.request_id, + data_size=data_size, + is_fast_path=is_fast, + ) + response_queue.put((identity, msgspec.msgpack.encode(resp))) + except Exception as e: + logger.error(f"Query handler error: {e}") + response_queue.put((identity, INFO_NOT_FOUND)) + + self._notify_listener(notify_addr) + + def _notify_listener(self, notify_addr: str) -> None: + try: + local = self._worker_local + sock = getattr(local, "notify_socket", None) + cached_addr = getattr(local, "notify_addr", None) + if sock is None or cached_addr != notify_addr: + if sock is not None: + sock.close(linger=0) + sock = self.zmq_ctx.socket(zmq.PUSH) + sock.connect(notify_addr) + local.notify_socket = sock + local.notify_addr = notify_addr + sock.send(b"", zmq.NOBLOCK) + except Exception: + local.notify_socket = None + local.notify_addr = None diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 6cf5c2f15..63466d116 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -29,6 +29,7 @@ class SharedMemoryConnector(OmniConnectorBase): self.stage_id = config.get("stage_id", -1) self.device = config.get("device", "cuda:0") self.threshold = int(config.get("shm_threshold_bytes", 65536)) + self.inline_small_payloads = bool(config.get("inline_small_payloads", False)) self._pending_keys: set[str] = set() self._metrics = { "puts": 0, @@ -53,9 +54,11 @@ class SharedMemoryConnector(OmniConnectorBase): payload = self.serialize_obj(data) size = len(payload) - # Currently, we always use SHM. - if True: - # Use Shared Memory + # The legacy async-chunk adapter transfers only the connector key + # across stages; it cannot deliver inline metadata to the receiver. + # Keep key-addressed SHM as the default and only inline payloads + # when the caller explicitly enables the metadata-aware path. + if size >= self.threshold or not self.inline_small_payloads: lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock" with open(lock_file, "wb+") as lockf: fcntl.flock(lockf, fcntl.LOCK_EX) @@ -67,9 +70,8 @@ class SharedMemoryConnector(OmniConnectorBase): self._pending_keys.add(put_key) self._metrics["shm_writes"] += 1 else: - # Inline - pass bytes directly to avoid double serialization of the object - # We already serialized it to check size, so we pass the bytes. - # The Queue will pickle these bytes (fast), avoiding re-serializing the complex object. + # Inline small payloads to avoid the mmap + file-lock overhead + # that dominates codec chunk transfers. metadata = {"inline_bytes": payload, "size": size} self._metrics["inline_writes"] += 1 diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index 7ee2dcde3..c49da1c6c 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/vllm_omni/distributed/omni_connectors/factory.py @@ -114,11 +114,23 @@ def _create_mooncake_transfer_engine_connector(config: dict[str, Any]) -> OmniCo return MooncakeTransferEngineConnector(config) +def _create_mori_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase: + try: + from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector + except ImportError: + import sys + + sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + from omni_connectors.connectors.mori_transfer_engine_connector import MoriTransferEngineConnector + return MoriTransferEngineConnector(config) + + # Register connectors OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_mooncake_store_connector) OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector) OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector) OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector) OmniConnectorFactory.register_connector("YuanrongTransferEngineConnector", _create_yuanrong_transfer_engine_connector) +OmniConnectorFactory.register_connector("MoriTransferEngineConnector", _create_mori_transfer_engine_connector) # Backward-compatible aliases – will be removed in the future OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector) diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 121fb79ec..eb1c3eec6 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -444,12 +444,37 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self, waiting_queue: Any, running_queue: list[Request], + *, + scheduler_requests: dict[str, Request] | None = None, ) -> None: """ Process pending chunks for waiting and running queues. + + When ``scheduler_requests`` is provided, purges any + ``waiting_for_chunk_*_requests`` deque entries whose + ``request_id`` is no longer tracked by it (e.g. after a + mid-flight abort that ran ``Scheduler._free_request``) before + processing chunks. Without this purge, ``restore_queues`` would + later re-inject the freed ``Request`` onto ``running_queue`` and + the worker's ``_update_states`` would crash with ``KeyError`` + reading ``self.requests[req_id]``. See vllm-project/vllm-omni#3736. + + ``scheduler_requests`` is keyword-only and optional; production + schedulers always pass their live request map, while legacy + callers that don't track aborts may omit it to keep the prior + (unguarded) behaviour. """ if self.connector.stage_id == 0: return + + # Purge deque entries whose request was freed mid-flight (abort → + # Scheduler._free_request) before any chunk processing, so neither + # the legacy nor the active-stream path can re-inject a zombie + # Request onto the queues. See vllm-project/vllm-omni#3736. + if scheduler_requests is not None: + self._purge_untracked_chunk_requests(self.waiting_for_chunk_waiting_requests, scheduler_requests) + self._purge_untracked_chunk_requests(self.waiting_for_chunk_running_requests, scheduler_requests) + if self._active_window <= 0: self._process_chunk_queue_legacy( waiting_queue, self.waiting_for_chunk_waiting_requests, RequestStatus.WAITING, self._finished_load_reqs @@ -558,6 +583,29 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.requests_origin_status[request.request_id] = target_status waiting_for_chunk_list.append(request) + def _purge_untracked_chunk_requests( + self, + deque_list: deque[Any], + scheduler_requests: dict[str, Request], + ) -> None: + """Drop deque entries whose ``request_id`` is not in + ``scheduler_requests`` and reclaim their receiver-side state. + + Handles requests that were aborted mid-flight while parked in a + chunk-transfer deque: ``Scheduler._free_request`` deleted the + entry from ``scheduler.requests`` but the deque still holds a + reference to the now-freed ``Request``. Order of survivors is + preserved. + """ + if not deque_list: + return + for _ in range(len(deque_list)): + request = deque_list.popleft() + if request.request_id in scheduler_requests: + deque_list.append(request) + else: + self.cleanup_receiver(request.request_id) + def restore_queues( self, waiting_queue: Any, @@ -566,7 +614,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): ) -> None: """ Restore requests waiting for chunk to the waiting and running queues. + + Re-runs the zombie purge first to close the race window where an + abort fires *between* ``process_pending_chunks`` and the + ``finally``-clause ``restore_queues`` call. Without the second + purge, ``running_queue.extend(...)`` would still re-inject a + freed ``Request`` and crash the worker on the next tick. + + ``scheduler_requests`` is optional for back-compat with legacy + callers (older tests pass only the two queue arguments). When + provided, it gates both the deque purge and the per-request + admit checks below; when ``None``, the purge is skipped and + every parked request is restored unconditionally (the + pre-purge behavior). """ + if scheduler_requests is not None: + self._purge_untracked_chunk_requests(self.waiting_for_chunk_waiting_requests, scheduler_requests) + self._purge_untracked_chunk_requests(self.waiting_for_chunk_running_requests, scheduler_requests) # Add request waiting for chunk to the waiting and running queue for request in self.waiting_for_chunk_waiting_requests: if scheduler_requests is None or request.request_id in scheduler_requests: diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index abe6e6a20..ee049904f 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -11,6 +11,7 @@ logger = get_connector_logger(__name__) TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset( { "MooncakeTransferEngineConnector", + "MoriTransferEngineConnector", "YuanrongTransferEngineConnector", } ) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 25244c462..148412012 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -35,6 +35,7 @@ def _register_omni_hf_configs() -> None: try: from transformers import AutoConfig + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingDenseConfig from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) @@ -55,6 +56,7 @@ def _register_omni_hf_configs() -> None: _CONFIG_REGISTRY = None for model_type, config_cls in [ + ("dense", MingDenseConfig), ("qwen3_tts", Qwen3TTSConfig), ("cosyvoice3", CosyVoice3Config), ("glm_tts", GLMTTSConfig), @@ -82,6 +84,9 @@ def register_omni_models_to_vllm(): if arch not in supported_archs: ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}") + # Register omni-specific reasoning parsers (e.g., step_audio). + import vllm_omni.reasoning # noqa: F401 + @dataclass class OmniEngineArgs(EngineArgs): diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index da1497c0c..05da3d3c3 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1363,6 +1363,7 @@ class AsyncOmniEngine: async def _run_orchestrator() -> None: self._initialize_stages(stage_init_timeout) + pd_config = self._detect_pd_config() coordinator_pub_address: str | None = None load_balancer_factory: Callable[[], LoadBalancer] | None = None diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 536f8d525..bb73f74a3 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -73,6 +73,8 @@ class UnregisterRemoteReplicaMessage(EngineQueueMessage, kw_only=True): class ErrorMessage(EngineQueueMessage, kw_only=True): type: Literal["error"] = "error" error: str + status_code: int | None = None + error_type: str | None = None fatal: bool = False request_id: str | None = None stage_id: int | None = None diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 017362812..f8ff84de8 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -689,7 +689,7 @@ class Orchestrator: await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) for eco in raw_outputs.outputs: req_state = self.request_states.get(getattr(eco, "request_id", None)) - if req_state is None: + if req_state is None or not req_state.streaming.enabled: continue req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False)) req_state.streaming.new_prompt_len_snapshot = getattr( @@ -697,6 +697,8 @@ class Orchestrator: "new_prompt_len_snapshot", None, ) + if req_state.streaming.enabled: + await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state) # OmniSchedulerMixin.make_stats() already throttles # per-scheduler at 1 Hz, so raw_outputs.scheduler_stats # being non-None means this replica passed its own gate. @@ -807,6 +809,8 @@ class Orchestrator: request_id=parent_id, stage_id=stage_id, error=output.error, + status_code=getattr(output, "error_status_code", None), + error_type=getattr(output, "error_type", None), ) ) await self._cleanup_request_ids( @@ -829,6 +833,34 @@ class Orchestrator: if self.request_states.pop(request_id, None) is not None and self._running_counter is not None: self._running_counter.decrement() + async def _apply_raw_terminal_stage_finish( + self, + stage_id: int, + eco: Any, + req_state: OrchestratorRequestState, + ) -> None: + """Record session-level finish markers dropped by the streaming output processor. + + Streaming segment stops set ``is_segment_finished=True`` and are handled + via processed outputs. Session termination (e.g. ``finish_requests`` after + ``resumable=False``) emits a terminal ``finish_reason`` with + ``is_segment_finished=False``, but vLLM's output processor may remove the + request state before that EngineCoreOutput is processed. + + Only update ``finished_final_output_stage_ids`` here. Request cleanup stays + in ``_route_output`` so downstream async-chunk stages can still deliver + outputs after stage-0 session end. + """ + if getattr(eco, "finish_reason", None) is None: + return + if getattr(eco, "is_segment_finished", False): + return + + final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} + if stage_id not in final_output_stage_ids: + return + req_state.finished_final_output_stage_ids.add(stage_id) + def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: """Attach CFG companion ids to diffusion sampling params when needed.""" companion_request_ids = self._cfg_tracker.get_companion_request_ids(request_id) @@ -872,7 +904,6 @@ class Orchestrator: req_state.finished_final_output_stage_ids.add(stage_id) final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids) - if self.stage_pools[stage_id].final_output: await self.output_async_queue.put( OutputMessage( diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 369fe3c7b..5c0b752b3 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -269,6 +269,10 @@ class OmniRequestState(RequestState): # tokens, so we always store a cumulative copy here. base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids) + # Attach cumulative_text only at the final step for inter-stage use. + if finish_reason is not None and hasattr(self.detokenizer, "output_text"): + base_output.cumulative_text = self.detokenizer.output_text + if not hasattr(base_output, "multimodal_output"): setattr(base_output, "multimodal_output", {}) if self.mm_accumulated: diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 32c92ff68..eb989db0e 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -89,6 +89,20 @@ class StageEngineCoreProc(EngineCoreProc): signal_callback: SignalCallback | None = None maybe_register_config_serialize_by_value() + # Register vllm-omni reasoning parsers (e.g. step_audio) in this + # subprocess so they are available when the engine core resolves + # ``--reasoning-parser``. The main process already registered them + # at import time, but the forked subprocess starts with a fresh + # ReasoningParserManager. + try: + import vllm_omni.reasoning # noqa: F401 + except ImportError: + logger.warning( + "Failed to import vllm_omni.reasoning in subprocess; " + "custom reasoning parsers (e.g. step_audio) will not be " + "available." + ) + engine_core: StageEngineCoreProc | None = None coord_client: OmniCoordClientForStage | None = None try: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 93f2dadff..6a1f505e4 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -33,6 +33,7 @@ from vllm_omni.entrypoints.omni_base import ( OmniBase, OmniEngineDeadError, ) +from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.outputs import OmniRequestOutput @@ -483,10 +484,13 @@ class AsyncOmni(EngineClient, OmniBase): except (asyncio.CancelledError, GeneratorExit): cancelled = True except Exception as error: + status_code, error_type = client_error_metadata(error) await req_state.queue.put( ErrorMessage( request_id=request_id, error=str(error), + status_code=status_code, + error_type=error_type, ) ) finally: @@ -589,7 +593,7 @@ class AsyncOmni(EngineClient, OmniBase): result.error, error_stage_id=result.stage_id, ) - raise RuntimeError(result.error) + self._raise_nonfatal_error_message(result) if not isinstance(result, OutputMessage): logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result) @@ -662,6 +666,18 @@ class AsyncOmni(EngineClient, OmniBase): await self.event_resolver.resolve(msg) continue + if isinstance(msg, ErrorMessage) and not msg.fatal: + req_state = self.request_states.get(msg.request_id) + if req_state is not None: + await req_state.queue.put(msg) + else: + logger.warning( + "[%s] dropping non-fatal error for unknown req %s", + self._name, + msg.request_id, + ) + continue + should_continue, _, stage_id, req_state = self._handle_output_message(msg) if should_continue: continue diff --git a/vllm_omni/entrypoints/chat_utils.py b/vllm_omni/entrypoints/chat_utils.py deleted file mode 100644 index 4c3d311ec..000000000 --- a/vllm_omni/entrypoints/chat_utils.py +++ /dev/null @@ -1,62 +0,0 @@ -import numpy as np - - -async def extract_audio_from_video_async(video_url: str) -> tuple[np.ndarray, int | float]: - """Extract audio from a video URL using vllm's load_audio. - - Returns a (audio_array, sample_rate) tuple compatible with audio format. - All blocking I/O operations are run in a thread pool. - """ - import asyncio - import os - import tempfile - from urllib.parse import urlparse - - parsed_url = urlparse(video_url) - temp_video_file_path = None - - def _download_video_sync(url: str) -> bytes: - from urllib.request import urlopen - - return urlopen(url).read() - - def _write_temp_file_sync(data: bytes, suffix: str) -> str: - with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file: - temp_file.write(data) - return temp_file.name - - def _load_audio_sync(file_path: str) -> tuple[np.ndarray, int | float]: - from vllm.multimodal.media.audio import load_audio - - return load_audio(file_path, sr=16000) - - def _cleanup_file_sync(file_path: str) -> None: - try: - if os.path.exists(file_path): - os.unlink(file_path) - except OSError: - pass - - try: - if parsed_url.scheme in ("http", "https"): - video_data = await asyncio.to_thread(_download_video_sync, video_url) - temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4") - elif parsed_url.scheme == "file": - from urllib.request import url2pathname - - temp_video_file_path = url2pathname(parsed_url.path) - elif parsed_url.scheme == "data": - import base64 - - header, data = video_url.split(",", 1) - video_data = base64.b64decode(data) - temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4") - else: - # Assume local file path - temp_video_file_path = video_url - - audio_array, sample_rate = await asyncio.to_thread(_load_audio_sync, temp_video_file_path) - return audio_array, sample_rate - finally: - if temp_video_file_path and parsed_url.scheme in ("http", "https", "data"): - await asyncio.to_thread(_cleanup_file_sync, temp_video_file_path) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index d415c07ea..029665f59 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -20,6 +20,7 @@ from vllm_omni.engine.messages import ( from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e +from vllm_omni.errors import raise_client_error_or from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_finalize from vllm_omni.metrics.prometheus import OmniPrometheusMetrics from vllm_omni.metrics.stats import OrchestratorAggregator @@ -372,7 +373,7 @@ class OmniBase(PDDisaggregationMixin): msg.error, error_stage_id=msg.stage_id, ) - raise RuntimeError(msg.error) + self._raise_nonfatal_error_message(msg) if not isinstance(msg, OutputMessage): logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg.type) @@ -408,6 +409,15 @@ class OmniBase(PDDisaggregationMixin): return False, req_id, stage_id, req_state + def _raise_nonfatal_error_message(self, msg: ErrorMessage) -> None: + """Raise the exception for a non-fatal, request-scoped error message.""" + raise_client_error_or( + msg.error, + status_code=msg.status_code, + error_type=msg.error_type, + fallback=RuntimeError, + ) + def _check_engine_output_error( self, result: OutputMessage, @@ -424,6 +434,8 @@ class OmniBase(PDDisaggregationMixin): error_text = getattr(engine_outputs, "error", None) if error_text is None: return + status_code = engine_outputs.error_status_code + error_type = engine_outputs.error_type logger.error( "[%s] Stage error for req=%s stage-%s: %s", self.__class__.__name__, @@ -437,7 +449,12 @@ class OmniBase(PDDisaggregationMixin): error_text, error_stage_id=stage_id, ) - raise EngineGenerateError(error_text) + raise_client_error_or( + error_text, + status_code=status_code, + error_type=error_type, + fallback=EngineGenerateError, + ) def _process_single_result( self, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py old mode 100644 new mode 100755 index 1be7aceba..c11db925f --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -91,6 +91,7 @@ from vllm.utils import random_uuid from vllm.utils.system_utils import decorate_logs from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError from vllm_omni.entrypoints.openai.image_api_utils import ( @@ -125,7 +126,12 @@ from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServin from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler -from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage +from vllm_omni.entrypoints.openai.serving_video import ( + OmniOpenAIServingVideo, + ReferenceAudio, + ReferenceImage, + ReferenceVideo, +) from vllm_omni.entrypoints.openai.serving_video_stream import OmniStreamingVideoHandler from vllm_omni.entrypoints.openai.stage_params import ( build_stage_sampling_params_list, @@ -134,8 +140,9 @@ from vllm_omni.entrypoints.openai.stage_params import ( from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request -from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference +from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url, decode_input_reference from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI +from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace @@ -424,6 +431,15 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + # In vllm-omni's multi-process architecture, create_engine_config() runs + # in the subprocess (StageEngineCoreProc), not in the API server process. + # Propagate reasoning_parser from CLI args into structured_outputs_config + # here so that OmniOpenAIServingChat receives the correct value. + # (Upstream vLLM does this in EngineArgs.create_engine_config().) + if hasattr(args, "reasoning_parser") and args.reasoning_parser: + if not args.structured_outputs_config.reasoning_parser: + args.structured_outputs_config.reasoning_parser = args.reasoning_parser + # Load logging config for uvicorn if specified log_config = get_uvicorn_log_config(args) if log_config is not None: @@ -1797,6 +1813,9 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise + except OmniClientError as e: + logger.info("Client error during image generation: %s", e) + raise HTTPException(status_code=e.status_code, detail=e.message) from e except ValueError as e: logger.error(f"Validation error: {e}") raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) @@ -2015,6 +2034,16 @@ async def edit_images( _update_if_not_none(gen_params, "layers", layers) _update_if_not_none(gen_params, "resolution", resolution) + extra_args = dict(getattr(gen_params, "extra_args", {}) or {}) + edit_extra_args = _build_hunyuan_edit_extra_args( + bot_task=bot_task, + sys_type=sys_type, + system_prompt=system_prompt, + ) + extra_args.update(edit_extra_args) + if extra_args: + gen_params.extra_args = extra_args + # 4. Generate images request_id = f"img_edit-{random_uuid()}" raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id) @@ -2142,6 +2171,9 @@ async def edit_images( return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise + except OmniClientError as e: + logger.info("Client error during image edit: %s", e) + raise HTTPException(status_code=e.status_code, detail=e.message) from e except ValueError as e: logger.error(f"Validation error: {e}") raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) @@ -2313,6 +2345,28 @@ def _check_max_generated_image_size( ) +def _build_hunyuan_edit_extra_args( + *, + bot_task: str | None, + sys_type: str | None, + system_prompt: str | None, +) -> dict[str, Any]: + """Map Hunyuan /v1/images/edits form fields to DiT ``extra_args``.""" + extra_args: dict[str, Any] = {} + effective_use_system_prompt = sys_type + if effective_use_system_prompt is None and bot_task is not None: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import resolve_sys_type + + effective_use_system_prompt = resolve_sys_type(bot_task) + if effective_use_system_prompt is not None: + extra_args["use_system_prompt"] = effective_use_system_prompt + if system_prompt is not None: + extra_args["system_prompt"] = system_prompt + if bot_task is not None: + extra_args["bot_task"] = bot_task + return extra_args + + def _update_if_not_none(object: Any, key: str, val: Any) -> None: if val is not None: setattr(object, key, val) @@ -2499,6 +2553,84 @@ def _parse_form_json(value: str | None, expected_type: type | None = None) -> An return parsed +def _config_get(config: Any, key: str, default: Any = None) -> Any: + if isinstance(config, dict): + return config.get(key, default) + if hasattr(config, "get"): + try: + return config.get(key, default) + except Exception: + pass + return getattr(config, key, default) + + +def _stage_engine_args(stage_cfg: Any) -> Any: + return _config_get(stage_cfg, "engine_args", {}) or {} + + +def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]: + if not stage_configs: + return [] + + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + model_classes: list[type] = [] + for stage_cfg in stage_configs: + if get_stage_type(stage_cfg) != "diffusion": + continue + model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name") + if not model_class_name: + continue + model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) + if model_cls is not None: + model_classes.append(model_cls) + return model_classes + + +def _normalize_reference_video_decode_spec(spec: ReferenceVideoDecodeSpec) -> ReferenceVideoDecodeSpec: + max_frames = spec.max_frames + if max_frames is not None: + try: + max_frames = int(max_frames) + except (TypeError, ValueError) as exc: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Invalid reference video decode spec: max_frames must be an integer.", + ) from exc + if max_frames <= 0: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Invalid reference video decode spec: max_frames must be positive.", + ) + + keep = str(spec.keep or "first").strip().lower() + if keep not in {"first", "last"}: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Invalid reference video decode spec: keep must be either 'first' or 'last'.", + ) + return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=cast(Literal["first", "last"], keep)) + + +def _reference_video_decode_spec( + req: VideoGenerationRequest, + stage_configs: list[Any] | None, +) -> ReferenceVideoDecodeSpec: + video_params = req.resolve_video_params() + extra_params = req.extra_params if isinstance(req.extra_params, dict) else {} + for model_cls in _diffusion_model_classes(stage_configs): + resolver = getattr(model_cls, "reference_video_decode_spec", None) + if resolver is None: + continue + try: + spec = resolver(num_frames=video_params.num_frames, extra_args=extra_params) + except (TypeError, ValueError) as exc: + raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc + if spec is not None: + return _normalize_reference_video_decode_spec(spec) + return ReferenceVideoDecodeSpec(max_frames=video_params.num_frames, keep="first") + + def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse: resp = VideoResponse( model=model_name, @@ -2542,6 +2674,9 @@ def _video_error_from_exception(exc: Exception) -> VideoError: message = str(exc.detail) if exc.detail else str(exc) return VideoError(code=exc.status_code, message=message) + if isinstance(exc, OmniClientError): + return VideoError(code=exc.status_code, message=exc.message) + if isinstance(exc, (EngineGenerateError, EngineDeadError)): err = create_error_response(exc) return VideoError(code=err.error.code, message=err.error.message) @@ -2565,6 +2700,8 @@ async def _run_video_generation_job( request: VideoGenerationRequest, video_id: str, reference_image: ReferenceImage | None = None, + reference_video: ReferenceVideo | None = None, + reference_audio: ReferenceAudio | None = None, app_state: Any | None = None, ) -> None: job = await VIDEO_STORE.get(video_id) @@ -2577,7 +2714,11 @@ async def _run_video_generation_job( output_path = None try: video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes( - request, video_id, reference_image=reference_image + request, + video_id, + reference_image=reference_image, + reference_video=reference_video, + reference_audio=reference_audio, ) file_name = f"{video_id}.{job.file_extension}" @@ -2634,9 +2775,12 @@ async def _run_video_generation_job( _cleanup_video(video_id, output_path) await VIDEO_STORE.pop(video_id) raise + finally: + if reference_audio is not None and os.path.exists(reference_audio.path): + os.unlink(reference_audio.path) -VIDEO_SYNC_TIMEOUT_S = 600.0 +VIDEO_SYNC_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_VIDEO_SYNC_TIMEOUT", 600.0)) async def _parse_video_form( @@ -2644,6 +2788,8 @@ async def _parse_video_form( prompt: str = Form(...), input_reference: UploadFile | None = File(default=None), image_reference: str | None = Form(default=None), + video_reference: str | None = Form(default=None), + audio_reference: str | None = Form(default=None), model: str | None = Form(default=None), seconds: SecondStr | None = Form(default=None), size: SizeStr | None = Form(default=None), @@ -2668,7 +2814,14 @@ async def _parse_video_form( frame_interpolation_model_path: str | None = Form(default=None), lora: str | None = Form(default=None), extra_params: str | None = Form(default=None), -) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None]: +) -> tuple[ + VideoGenerationRequest, + "OmniOpenAIServingVideo", + str, + ReferenceImage | None, + ReferenceVideo | None, + ReferenceAudio | None, +]: """FastAPI dependency that parses video form data, validates inputs, resolves the handler, and decodes any reference image. @@ -2676,11 +2829,16 @@ async def _parse_video_form( """ input_reference_bytes = await input_reference.read() if input_reference is not None else None parsed_image_reference = _parse_form_json(image_reference) + parsed_video_reference = _parse_form_json(video_reference) + parsed_audio_reference = _parse_form_json(audio_reference) - if parsed_image_reference is not None and input_reference_bytes is not None: + provided_references = sum( + item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference_bytes) + ) + if provided_references > 1: raise HTTPException( status_code=HTTPStatus.BAD_REQUEST.value, - detail="Provide either input_reference or image_reference, not both.", + detail="Provide only one of input_reference, image_reference, or video_reference.", ) request_data: dict[str, Any] = { @@ -2689,6 +2847,8 @@ async def _parse_video_form( "seconds": seconds, "size": size, "image_reference": parsed_image_reference, + "video_reference": parsed_video_reference, + "audio_reference": parsed_audio_reference, "user": user, "width": width, "height": height, @@ -2742,13 +2902,37 @@ async def _parse_video_form( detail=f"Video generation setup failed: {str(e)}", ) + decode_spec = ReferenceVideoDecodeSpec() + if parsed_video_reference is not None or input_reference_bytes is not None: + stage_configs = ( + handler.stage_configs + or app_stage_configs + or getattr(getattr(handler, "_engine_client", None), "stage_configs", None) + ) + decode_spec = _reference_video_decode_spec(request, stage_configs) try: - image_data = await decode_input_reference(request.image_reference, input_reference_bytes) + media_data = await decode_input_reference( + request.image_reference, + request.video_reference, + input_reference_bytes, + max_video_frames=decode_spec.max_frames, + video_keep=decode_spec.keep, + ) except InvalidInputReferenceError as exc: raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc - reference_image = ReferenceImage(data=image_data) if image_data is not None else None - return request, handler, effective_model_name, reference_image + reference_image = ReferenceImage(data=media_data) if isinstance(media_data, Image.Image) else None + reference_video = ReferenceVideo(data=media_data) if isinstance(media_data, list) else None + + reference_audio: ReferenceAudio | None = None + if request.audio_reference is not None: + try: + audio_path = await decode_audio_url(request.audio_reference.audio_url) + except InvalidInputReferenceError as exc: + raise HTTPException(400, detail=str(exc)) from exc + reference_audio = ReferenceAudio(path=audio_path) + + return request, handler, effective_model_name, reference_image, reference_video, reference_audio @router.post( @@ -2762,18 +2946,33 @@ async def _parse_video_form( ) async def create_video( raw_request: Request, - ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form), + ctx: tuple[ + VideoGenerationRequest, + OmniOpenAIServingVideo, + str, + ReferenceImage | None, + ReferenceVideo | None, + ReferenceAudio | None, + ] = Depends(_parse_video_form), ) -> VideoResponse: """Create an asynchronous video generation job. Accepts multipart form-data (see ``_parse_video_form`` for parameters), persists a queued job record, and starts generation in the background. """ - request, handler, effective_model_name, reference_image = ctx + request, handler, effective_model_name, reference_image, reference_video, reference_audio = ctx ref = video_response_from_request(effective_model_name, request) await VIDEO_STORE.upsert(ref.id, ref) task = asyncio.create_task( - _run_video_generation_job(handler, request, ref.id, reference_image, app_state=raw_request.app.state) + _run_video_generation_job( + handler, + request, + ref.id, + reference_image, + reference_video, + reference_audio, + app_state=raw_request.app.state, + ) ) await VIDEO_TASKS.upsert(ref.id, task) return ref @@ -2790,7 +2989,14 @@ async def create_video( ) async def create_video_sync( raw_request: Request, - ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form), + ctx: tuple[ + VideoGenerationRequest, + OmniOpenAIServingVideo, + str, + ReferenceImage | None, + ReferenceVideo | None, + ReferenceAudio | None, + ] = Depends(_parse_video_form), ) -> Response: """Synchronous video generation endpoint. @@ -2801,13 +3007,19 @@ async def create_video_sync( Metadata is returned via response headers ``X-Request-Id``, ``X-Model``, and ``X-Inference-Time-S``. """ - request, handler, effective_model_name, reference_image = ctx + request, handler, effective_model_name, reference_image, reference_video, reference_audio = ctx request_id = f"video_sync-{random_uuid()}" raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id) started_at = time.perf_counter() try: video_bytes, stage_durations, peak_memory_mb, _action = await asyncio.wait_for( - handler.generate_video_bytes(request, request_id, reference_image=reference_image), + handler.generate_video_bytes( + request, + request_id, + reference_image=reference_image, + reference_video=reference_video, + reference_audio=reference_audio, + ), timeout=VIDEO_SYNC_TIMEOUT_S, ) except asyncio.TimeoutError: @@ -2819,12 +3031,18 @@ async def create_video_sync( return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise + except OmniClientError as exc: + logger.info("Client error during sync video generation: %s", exc) + raise HTTPException(status_code=exc.status_code, detail=exc.message) from exc except Exception as exc: logger.exception("Sync video generation failed for request_id=%s", request_id) raise HTTPException( status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Video generation failed: {str(exc)}", ) from exc + finally: + if reference_audio is not None and os.path.exists(reference_audio.path): + os.unlink(reference_audio.path) inference_time_s = time.perf_counter() - started_at return Response( diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 9a858868e..c7ad1c9ef 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -7,6 +7,43 @@ from pydantic import AliasChoices, BaseModel, Field, field_validator, model_vali _MAX_EMBEDDING_DIM = 8192 +def _normalize_ref_audio_value(value): + if value is None: + return None + if isinstance(value, str): + return value + if isinstance(value, (list, tuple)): + items = [] + for item in value: + if not isinstance(item, str): + raise TypeError("'ref_audio' list entries must be strings") + items.append(item) + if not items: + raise ValueError("'ref_audio' list cannot be empty") + return items + raise TypeError("'ref_audio' must be a string or list of strings") + + +def _normalize_speaker_embedding_value(value): + if value is None: + return None + if not isinstance(value, (list, tuple)): + raise TypeError("'speaker_embedding' must be a list of numbers or list of embedding vectors") + if not value: + return [] + + first = value[0] + if isinstance(first, (list, tuple)): + embeddings = [] + for item in value: + if not isinstance(item, (list, tuple)): + raise TypeError("'speaker_embedding' must not mix flat and nested values") + embeddings.append([float(x) for x in item]) + return embeddings + + return [float(x) for x in value] + + class OpenAICreateSpeechRequest(BaseModel): input: str model: str | None = None @@ -46,7 +83,7 @@ class OpenAICreateSpeechRequest(BaseModel): default=None, description="Language code (e.g., 'Chinese', 'English', 'Auto')", ) - ref_audio: str | None = Field( + ref_audio: str | list[str] | None = Field( default=None, description="Reference audio for voice cloning (Base task). URL, base64, or file URI.", ) @@ -73,7 +110,7 @@ class OpenAICreateSpeechRequest(BaseModel): default=None, description="Use speaker embedding only without in-context learning (Base task)", ) - speaker_embedding: list[float] | None = Field( + speaker_embedding: list[float] | list[list[float]] | None = Field( default=None, max_length=_MAX_EMBEDDING_DIM, description="Pre-computed speaker embedding vector (1024-dim for 0.6B, " @@ -96,6 +133,14 @@ class OpenAICreateSpeechRequest(BaseModel): ge=0, description="Per-request initial chunk size override. If null, computed dynamically based on server load.", ) + non_streaming_mode: bool | None = Field( + default=None, + description=( + "Qwen3-TTS prompt construction mode override. This does not " + "control HTTP response streaming or async-chunk pipelining. " + "When null, use model defaults: Base=False, CustomVoice/VoiceDesign=True." + ), + ) extra_params: dict[str, Any] | None = Field( default=None, description=("Optional model-specific parameters passed directly to the model's extra_args."), @@ -108,13 +153,80 @@ class OpenAICreateSpeechRequest(BaseModel): raise ValueError("'sse' is not a supported stream_format yet. Please use 'audio'.") return v + @field_validator("ref_audio", mode="before") + @classmethod + def normalize_ref_audio(cls, v): + return _normalize_ref_audio_value(v) + @field_validator("speaker_embedding") @classmethod - def validate_speaker_embedding(cls, v: list[float] | None) -> list[float] | None: - if v is not None and not all(math.isfinite(x) for x in v): + def validate_speaker_embedding( + cls, v: list[float] | list[list[float]] | None + ) -> list[float] | list[list[float]] | None: + v = _normalize_speaker_embedding_value(v) + if v is None: + return None + if not v: + return [] + if isinstance(v[0], list): + for item in v: + if not item: + raise ValueError("'speaker_embedding' nested vectors must be non-empty") + if not all(math.isfinite(x) for x in item): + raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)") + return v + if not all(math.isfinite(x) for x in v): raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)") return v + @model_validator(mode="before") + @classmethod + def normalize_references_alias(cls, data): + """Map the BosonAI Higgs Audio v3 cookbook ``references`` array onto ``ref_audio`` / ``ref_text``. + + Upstream Higgs Audio v3 examples and the boson.ai cookbook send voice + clones as:: + + {"references": [{"audio_path": str, "text": str}]} + + vllm-omni's existing convention is ``ref_audio`` + ``ref_text``. + Without this alias pydantic silently drops the unknown ``references`` + field and the request falls through to zero-shot synthesis with no + error, which is hard to debug from the client side. The normalizer + translates a single reference and rejects multi-reference payloads + (vllm-omni does not yet support multi-shot voice clone) or conflicts + with the explicit ``ref_audio`` / ``ref_text`` fields. + """ + if not isinstance(data, dict): + return data + refs = data.get("references") + if refs is None: + return data + if not isinstance(refs, list) or len(refs) == 0: + raise ValueError("'references' must be a non-empty array of {audio_path, text} objects") + if len(refs) > 1: + raise ValueError("'references' only supports a single reference; multi-shot voice clone is not supported") + first = refs[0] + if not isinstance(first, dict): + raise ValueError("'references[0]' must be an object with 'audio_path' and optional 'text'") + audio_path = first.get("audio_path") + if not isinstance(audio_path, str) or not audio_path: + raise ValueError( + "'references[0].audio_path' is required and must be a string (URL, data: URI, or file path)" + ) + existing_ref_audio = data.get("ref_audio") + if existing_ref_audio and existing_ref_audio != audio_path: + raise ValueError("'references' and 'ref_audio' are mutually exclusive") + data["ref_audio"] = audio_path + text = first.get("text") + if isinstance(text, str) and text.strip(): + existing_ref_text = data.get("ref_text") + if existing_ref_text and existing_ref_text != text: + raise ValueError("'references[0].text' and 'ref_text' conflict; supply only one") + data["ref_text"] = text + data.pop("references", None) + return data + @model_validator(mode="before") @classmethod def reject_higgs_audio_v2_unsupported_aliases(cls, data): @@ -162,7 +274,7 @@ class OpenAICreateSpeechRequest(BaseModel): @model_validator(mode="after") def validate_embedding_constraints(self) -> "OpenAICreateSpeechRequest": if self.speaker_embedding is not None: - if self.ref_audio is not None: + if self.ref_audio is not None and not isinstance(self.ref_audio, list): raise ValueError("'speaker_embedding' and 'ref_audio' are mutually exclusive") return self @@ -266,6 +378,7 @@ class SpeechBatchItem(BaseModel): x_vector_only_mode: bool | None = None max_new_tokens: int | None = None initial_codec_chunk_frames: int | None = Field(default=None, ge=0) + non_streaming_mode: bool | None = None class BatchSpeechRequest(BaseModel): @@ -285,6 +398,7 @@ class BatchSpeechRequest(BaseModel): x_vector_only_mode: bool | None = None max_new_tokens: int | None = None initial_codec_chunk_frames: int | None = Field(default=None, ge=0) + non_streaming_mode: bool | None = None class SpeechBatchItemResult(BaseModel): @@ -319,6 +433,14 @@ class StreamingSpeechSessionConfig(BaseModel): ge=0, description="Initial chunk size for reduced TTFA. Overrides stage config for this session.", ) + non_streaming_mode: bool | None = Field( + default=None, + description=( + "Qwen3-TTS prompt construction mode override. This does not " + "control WebSocket audio streaming or async-chunk pipelining. " + "When null, use model defaults: Base=False, CustomVoice/VoiceDesign=True." + ), + ) ref_audio: str | None = None ref_text: str | None = None x_vector_only_mode: bool | None = None diff --git a/vllm_omni/entrypoints/openai/protocol/videos.py b/vllm_omni/entrypoints/openai/protocol/videos.py index ec5ab14e8..8773dd2d8 100644 --- a/vllm_omni/entrypoints/openai/protocol/videos.py +++ b/vllm_omni/entrypoints/openai/protocol/videos.py @@ -73,6 +73,27 @@ class UrlImageReference(BaseModel): ImageReference = UrlImageReference | FileImageReference +class FileVideoReference(BaseModel): + model_config = ConfigDict(extra="forbid") + file_id: str + + +class UrlVideoReference(BaseModel): + model_config = ConfigDict(extra="forbid") + video_url: str + + +VideoReference = UrlVideoReference | FileVideoReference + + +class UrlAudioReference(BaseModel): + model_config = ConfigDict(extra="forbid") + audio_url: str + + +AudioReference = UrlAudioReference + + class VideoGenerationRequest(BaseModel): """ OpenAI-style video generation request. @@ -99,6 +120,14 @@ class VideoGenerationRequest(BaseModel): default=None, description="Optional JSON-safe image reference that guides generation. Provide either image_url or file_id.", ) + video_reference: VideoReference | None = Field( + default=None, + description="Optional JSON-safe video reference that guides generation. Provide either video_url or file_id.", + ) + audio_reference: AudioReference | None = Field( + default=None, + description="Optional audio reference for speech-to-video. Provide audio_url (http(s) or data URL).", + ) # Video params block for extensibility video_params: VideoParams | None = Field(default=None, description="Optional video-specific parameters") diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index 1bc3df64e..5cf936e48 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -86,9 +86,14 @@ class RealtimeConnection(VllmRealtimeConnection): return [], 24000 sr = mm.get("sr") or mm.get("sample_rate") or mm.get("audio_sample_rate") or 24000 + if isinstance(sr, (list, tuple)) and sr: + sr = sr[-1] + if hasattr(sr, "item"): + sr = sr.item() + sample_rate_hz = int(sr) key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None) if key is None: - return [], int(sr) + return [], sample_rate_hz raw_audio = mm.get(key) chunks: list[np.ndarray] = [] @@ -101,7 +106,7 @@ class RealtimeConnection(VllmRealtimeConnection): arr = self._tensor_to_numpy(raw_audio) if arr is not None and arr.size > 0: chunks.extend(self._raw_waveform_to_deltas(arr)) - return chunks, int(sr) + return chunks, sample_rate_hz @staticmethod def _pcm16_b64(audio_f32: np.ndarray) -> str: diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py old mode 100644 new mode 100755 index c7b24e11f..35cfafcff --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -12,8 +12,16 @@ from typing import Any, Final, cast import jinja2 import torch from fastapi import Request +from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio from PIL import Image from pydantic import TypeAdapter +from vllm.entrypoints.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, + ConversationMessage, + get_history_tool_calls_cnt, + make_tool_call_id, +) from vllm_omni.diffusion.diffusion_engine import get_extra_body_params, get_extra_output_params from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -32,14 +40,6 @@ except ImportError: soundfile = None -from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, - ConversationMessage, - get_history_tool_calls_cnt, - make_tool_call_id, -) from vllm.entrypoints.launcher import terminate_if_errored from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, @@ -113,6 +113,7 @@ from vllm_omni.entrypoints.openai.utils import ( resolve_diffusion_od_config, validate_requested_speaker, ) +from vllm_omni.errors import OmniClientError from vllm_omni.lora.request import LoRARequest from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate @@ -652,19 +653,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) tok_params = request.build_tok_params(self.model_config) + mm_config = self.model_config.multimodal_config chat_params = request.build_chat_params( default_template, default_template_content_format, - ).with_defaults(default_template_kwargs) - - # OMNI: When use_audio_in_video=True, the qwen2_5_omni_thinker mm - # processor asserts that audio items are present alongside video items - # (inside render_chat_async). We must inject audio_url items into the - # messages BEFORE calling render_chat_async so the mm processor can - # count them correctly during tokenisation. - mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {} - if mm_proc_kw.get("use_audio_in_video", False): - messages = await self._inject_audio_from_video_urls(messages) + ).with_defaults( + default_template_kwargs, + default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None), + default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), + ) (conversation,), (engine_prompt,) = await renderer.render_chat_async( [messages], @@ -735,95 +732,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return conversation, [engine_prompt] - async def _inject_audio_from_video_urls( - self, - messages: list[ChatCompletionMessageParam], - ) -> list[ChatCompletionMessageParam]: - """Pre-extract audio from video URLs and inject as audio_url content items. - - When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal - processor requires that the number of audio items equals the number of - video items (it subtracts mm_counts["video"] from mm_counts["audio"]). - The client only sends video_url items; this method adds the matching - audio_url items on the server side before the renderer processes them. - """ - import io - - from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async - - new_messages: list[ChatCompletionMessageParam] = [] - for msg in messages: - content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None) - if not isinstance(content, list): - new_messages.append(msg) - continue - - video_urls = [ - part.get("video_url", {}).get("url") - for part in content - if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url") - ] - - if not video_urls: - new_messages.append(msg) - continue - - audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls)) - - audio_items: list[dict] = [] - for audio_array, sample_rate in audios: - buf = io.BytesIO() - if soundfile is not None: - soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV") - else: - import struct - - import numpy as np - - audio_np = np.asarray(audio_array, dtype=np.float32) - sr = int(sample_rate) - num_channels = 1 - bits_per_sample = 32 - num_frames = len(audio_np) - data_size = num_frames * num_channels * (bits_per_sample // 8) - # Write minimal RIFF/WAV header - buf.write(b"RIFF") - buf.write(struct.pack(" list[Any]: """Convert request dicts to stage-typed sampling params objects. @@ -2599,6 +2507,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): mm_processor_kwargs["target_h"] = height if width is not None: mm_processor_kwargs["target_w"] = width + if seed is not None and engine_prompt_data is not None: + mm_processor_kwargs["vae_generator_seed"] = int(seed) if mm_processor_kwargs: engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs if engine_prompt_data is not None: @@ -2960,6 +2870,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "[OmniOpenAIServingChat] Engine dead during streaming image edit, " "but cannot terminate server (no raw_request context)." ) + except OmniClientError as exc: + logger.info("Client error during streaming image edit: %s", exc) + chunk = ImageEditStreamError( + created=created, + model=model, + error={ + "message": exc.message, + "type": exc.error_type, + "code": exc.status_code, + }, + ) + yield f"data: {chunk.model_dump_json()}\n\n" except Exception as exc: logger.exception("Streaming image edit failed: %s", exc) chunk = ImageEditStreamError( @@ -3356,6 +3278,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return response + except OmniClientError as e: + logger.info("Client error during diffusion chat completion: %s", e) + return self._create_error_response( + e.message, + err_type=e.error_type, + status_code=e.status_code, + ) except Exception as e: logger.exception("Diffusion chat completion failed: %s", e) return self._create_error_response( diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 57bed35eb..99df2332b 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -9,6 +9,7 @@ import re import struct import time from collections import OrderedDict +from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from http import HTTPStatus from pathlib import Path @@ -65,6 +66,7 @@ from vllm_omni.utils.speaker_cache import ( logger = init_logger(__name__) # TTS Configuration +_MING_TTS_MODEL_ARCHS = {"MingTTSForConditionalGeneration"} _VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"} _QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"} _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} @@ -76,6 +78,7 @@ _MING_TTS_MODEL_STAGES = {"ming_tts"} _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"} _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"} +_HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"} _GLM_TTS_MODEL_STAGES = {"glm_tts"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES @@ -84,6 +87,7 @@ _TTS_MODEL_STAGES: set[str] = ( | _COSYVOICE3_TTS_MODEL_STAGES | _OMNIVOICE_TTS_MODEL_STAGES | _HIGGS_AUDIO_V2_TTS_MODEL_STAGES + | _HIGGS_V3_TTS_MODEL_STAGES | _COVO_AUDIO_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES | _MING_TTS_MODEL_STAGES @@ -98,28 +102,34 @@ _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = { "cosyvoice3", "voxcpm2", "higgs_audio_v2", + "higgs_audio_v3", } -_TTS_LANGUAGES: set[str] = { - "Auto", - "Chinese", - "English", - "Japanese", - "Korean", - "German", - "French", - "Russian", - "Portuguese", - "Spanish", - "Italian", -} +_TTS_LANGUAGES = frozenset( + { + "Auto", + "Chinese", + "English", + "Japanese", + "Korean", + "German", + "French", + "Russian", + "Portuguese", + "Spanish", + "Italian", + } +) _REF_AUDIO_MIN_DURATION = 1.0 # seconds _REF_AUDIO_MAX_DURATION = 30.0 # seconds _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256 _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 256 * 1024 * 1024 +_HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES = 256 +_HIGGS_V3_REF_CODE_CACHE_MAX_BYTES = 64 * 1024 * 1024 _QWEN3_TTS_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" _TTS_MAX_INSTRUCTIONS_LENGTH = 500 _TTS_MAX_NEW_TOKENS_MIN = 1 _TTS_MAX_NEW_TOKENS_MAX = 4096 +_MING_DEFAULT_PROMPT = MING_DEFAULT_PROMPT def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: int = 16) -> bytes: @@ -280,6 +290,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES self._ref_audio_model_artifact_ready: set[str] = set() self._request_ref_audio_artifact_keys: dict[str, str] = {} + self._higgs_audio_v3_ref_code_cache: OrderedDict[str, tuple[torch.Tensor, int]] = OrderedDict() + self._higgs_audio_v3_ref_code_cache_bytes = 0 + self._higgs_audio_v3_ref_code_inflight: dict[str, asyncio.Task[torch.Tensor]] = {} self._speaker_cache = get_speaker_cache() self._last_upload_ts = 0 self._upload_lock = asyncio.Lock() @@ -404,6 +417,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) self._fish_speech_tokenizer = None self._covo_audio_tokenizer = None + # Cached per process: the CosyVoice3 Qwen tokenizer + resolved model + # path used for dynamic-token sizing. Without this, every request + # re-ran snapshot_download + reloaded the tokenizer (~100 ms on the + # TTFP critical path) in _apply_cosyvoice3_dynamic_tokens. + self._cosyvoice3_tokenizer = None self._is_cosyvoice3 = ( self._tts_stage is not None @@ -431,6 +449,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Merge built-in speakers into the set initialized by _init_speaker_storage. self.supported_speakers |= self._load_supported_speakers() self.supported_speakers |= set(self.precomputed_speakers) + + self.supported_languages = self._load_supported_languages() + self._tts_tokenizer = None self._voxcpm2_tokenizer = None self._voxcpm2_split_map: dict[int, list[int]] = {} @@ -514,6 +535,32 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _load_codec_frame_rate(self) -> float | None: """Load codec frame rate from speech tokenizer config for prompt length estimation.""" + if self._tts_model_type == "ming_tts": + try: + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingTTSConfig + + hf_config = self.engine_client.model_config.hf_config + ming_cfg = MingTTSConfig.from_hf_config(hf_config) + patch_size = int(ming_cfg.patch_size) + audio_frame_hop = int(ming_cfg.audio_frame_hop) + sample_rate = int(ming_cfg.sample_rate) + if patch_size <= 0 or audio_frame_hop <= 0 or sample_rate <= 0: + raise ValueError( + "Ming config has invalid tokenizer timing values: " + f"patch_size={patch_size}, audio_frame_hop={audio_frame_hop}, sample_rate={sample_rate}" + ) + rate = float(sample_rate) / float(audio_frame_hop * patch_size) + logger.info( + "Derived Ming codec frame rate: %.1f Hz (sample_rate=%s, audio_frame_hop=%s, patch_size=%s)", + rate, + sample_rate, + audio_frame_hop, + patch_size, + ) + return rate + except Exception as e: + logger.warning(f"Failed to derive Ming codec frame rate from hf_config: {e}") + try: model_path = self.engine_client.model_config.model st_config_path = os.path.join(model_path, "speech_tokenizer", "config.json") @@ -558,7 +605,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _find_tts_stage(self): """Find and return the TTS stage config, or None if not found.""" for stage in self.engine_client.stage_configs: - if stage.engine_args.model_stage in _TTS_MODEL_STAGES: + engine_args = stage.engine_args + model_stage = engine_args.model_stage + model_arch = getattr(engine_args, "model_arch", None) + worker_type = getattr(engine_args, "worker_type", None) + if model_stage in _TTS_MODEL_STAGES: + return stage + # Ming dense identifies its AR entry stage by architecture because + # it does not use a dedicated TTS model_stage value. + if model_arch in _MING_TTS_MODEL_ARCHS and worker_type == "ar": return stage return None @@ -588,12 +643,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "voxcpm2" if model_stage in _MING_TTS_MODEL_STAGES: return "ming_flash_omni_tts" + if model_arch in _MING_TTS_MODEL_ARCHS: + return "ming_tts" if model_stage in _MOSS_TTS_MODEL_STAGES: return "moss_tts_nano" if model_stage in _MOSS_TTS_FULL_MODEL_STAGES: return "moss_tts" if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES: return "higgs_audio_v2" + if model_stage in _HIGGS_V3_TTS_MODEL_STAGES: + return "higgs_audio_v3" if model_stage in _GLM_TTS_MODEL_STAGES: return "glm_tts" return None @@ -666,6 +725,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): try: if self._tts_model_type == "glm_tts": return set() + if self._tts_model_type == "ming_tts": + return set() if self._tts_model_type == "voxcpm2": return {"default"} if self._tts_model_type == "voxtral_tts": @@ -689,6 +750,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return set() + def _load_supported_languages(self) -> frozenset[str]: + """Load supported languages (title-cased) from the model configuration""" + if self._tts_model_type != "qwen3_tts": + return _TTS_LANGUAGES + try: + config = self.engine_client.model_config.hf_config.talker_config + + if isinstance(config, dict): + codec_language_id = config.get("codec_language_id") + else: + codec_language_id = getattr(config, "codec_language_id", None) + + if codec_language_id and isinstance(codec_language_id, Mapping): + return frozenset(str(language).title() for language in codec_language_id) | {"Auto"} + + logger.warning("No codec_language_id found in talker_config; falling back to default languages") + except Exception as e: + logger.warning("Could not load languages from model config: %s", e) + return _TTS_LANGUAGES + def _estimate_ref_code_len(self, ref_audio: object) -> int | None: """Estimate ref_code length from ref_audio waveform without running the codec. @@ -934,6 +1015,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "moss_tts_nano", "glm_tts", "higgs_audio_v2", + "higgs_audio_v3", ): label = { "cosyvoice3": "CosyVoice3", @@ -941,6 +1023,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "omnivoice": "OmniVoice", "moss_tts_nano": "MOSS-TTS-Nano", "higgs_audio_v2": "Higgs-Audio V2", + "higgs_audio_v3": "Higgs-Audio V3", "glm_tts": "GLM-TTS", }.get(self._tts_model_type, self._tts_model_type) return ( @@ -1176,9 +1259,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)") emb_dim = len(embedding) - dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim) - if dim_err is not None: - raise ValueError(dim_err) + if self._tts_model_type == "ming_tts": + if emb_dim != 192: + logger.warning("speaker_embedding has %d dimensions; Ming dense expects 192", emb_dim) + else: + dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim) + if dim_err is not None: + raise ValueError(dim_err) async with self._upload_lock: voice_name_lower = name.lower() @@ -1265,7 +1352,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _is_tts_model(self) -> bool: """Check if the current model is a supported TTS model.""" - return any(stage.engine_args.model_stage in _TTS_MODEL_STAGES for stage in self.engine_client.stage_configs) + return self._find_tts_stage() is not None def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate TTS request parameters. Returns error message or None.""" @@ -1278,6 +1365,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if self._tts_model_type == "voxcpm2": return self._validate_voxcpm2_request(request) if self._tts_model_type == "ming_flash_omni_tts": + return self._validate_ming_flash_omni_tts_request(request) + if self._tts_model_type == "ming_tts": return self._validate_ming_tts_request(request) if self._tts_model_type in ("moss_tts_nano", "moss_tts"): return self._validate_moss_tts_request(request) @@ -1324,7 +1413,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ids = self._voxcpm2_tokenizer.encode(text, add_special_tokens=True) return split_multichar_chinese(ids, self._voxcpm2_split_map) - def _validate_ming_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + def _validate_ming_flash_omni_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate Ming-flash-omni standalone-talker request parameters.""" if not request.input or not request.input.strip(): return "Input text cannot be empty" @@ -1361,6 +1450,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _validate_ref_audio_format(self, ref_audio: str) -> str | None: """Validate ref_audio is a supported URI format. Returns error or None.""" + if not isinstance(ref_audio, str): + return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" if not ( ref_audio.startswith(("http://", "https://")) or ref_audio.startswith("data:") @@ -1413,9 +1504,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if not request.input or not request.input.strip(): return "Input text cannot be empty" - # Validate language - if request.language is not None and request.language not in _TTS_LANGUAGES: - return f"Invalid language '{request.language}'. Supported: {', '.join(sorted(_TTS_LANGUAGES))}" + # Validate language (case-insensitive; normalized to the title-cased config form) + if request.language is not None: + request.language = request.language.title() + if request.language not in self.supported_languages: + return ( + f"Invalid language '{request.language}'. Supported: {', '.join(sorted(self.supported_languages))}" + ) # Validate speaker for CustomVoice task if task_type == "CustomVoice": @@ -1896,6 +1991,153 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._higgs_audio_v2_processor = processor return processor + # ---- higgs-audio v3 ---- + + def _validate_higgs_audio_v3_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate higgs_audio_v3 request parameters.""" + if not request.input or not request.input.strip(): + return "higgs_audio_v3: input text cannot be empty" + if request.ref_audio is not None and not request.ref_text: + # Voice clone ref_text is optional for v3 (improves fidelity but not required) + pass + if request.max_new_tokens is not None: + if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: + return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" + return None + + async def _build_higgs_audio_v3_params(self, request: OpenAICreateSpeechRequest): + """Build prompt_token_ids for higgs_audio_v3. + + Plain-text path: builds ``[tts, text, tokens, audio]``. + Voice-clone path: encodes reference audio, applies delay pattern, + builds ``[tts, (ref_text, tokens,) ref_audio, -100xN, text, tokens, audio]``. + """ + adapter = await self._resolve_higgs_audio_v3_adapter() + + if request.ref_audio is None: + prompt_ids = adapter.build_prompt(request.input) + return tokens_input(prompt_token_ids=prompt_ids) + + # Voice clone + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + apply_delay_pattern, + encode_reference_audio, + ) + + wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + artifact_key = self._get_resolved_ref_audio_artifact_key(request.ref_audio) + wav = np.asarray(wav_list, dtype=np.float32) + ref_codes_delayed, cache_hit, inflight_wait = await self._resolve_higgs_audio_v3_ref_codes( + artifact_key, + wav, + int(sr), + encode_reference_audio, + apply_delay_pattern, + ) + del cache_hit, inflight_wait + + prompt_ids = adapter.build_prompt( + request.input, + num_ref_tokens=int(ref_codes_delayed.shape[0]), + reference_text=request.ref_text or None, + ) + prompt = tokens_input(prompt_token_ids=prompt_ids) + import torch + + prompt["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), + } + return prompt + + async def _resolve_higgs_audio_v3_ref_codes( + self, + artifact_key: str | None, + wav: np.ndarray, + sr: int, + encode_reference_audio, + apply_delay_pattern, + ) -> tuple[torch.Tensor, bool, bool]: + ref_codes_delayed = self._get_higgs_audio_v3_ref_codes(artifact_key) + if ref_codes_delayed is not None: + return ref_codes_delayed, True, False + if not artifact_key: + ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) + return apply_delay_pattern(ref_codes_raw), False, False + + task = self._higgs_audio_v3_ref_code_inflight.get(artifact_key) + if task is not None: + return (await task).clone(), False, True + + async def _encode_and_cache() -> torch.Tensor: + ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) + delayed = apply_delay_pattern(ref_codes_raw) + self._put_higgs_audio_v3_ref_codes(artifact_key, delayed) + cached = self._get_higgs_audio_v3_ref_codes(artifact_key) + return cached if cached is not None else delayed.detach().to("cpu", dtype=torch.long).contiguous() + + task = asyncio.create_task(_encode_and_cache()) + self._higgs_audio_v3_ref_code_inflight[artifact_key] = task + try: + return (await task).clone(), False, False + finally: + if self._higgs_audio_v3_ref_code_inflight.get(artifact_key) is task: + self._higgs_audio_v3_ref_code_inflight.pop(artifact_key, None) + + def _get_higgs_audio_v3_ref_codes(self, artifact_key: str | None) -> torch.Tensor | None: + if not artifact_key: + return None + cached = self._higgs_audio_v3_ref_code_cache.get(artifact_key) + if cached is None: + return None + self._higgs_audio_v3_ref_code_cache.move_to_end(artifact_key) + return cached[0].clone() + + def _put_higgs_audio_v3_ref_codes(self, artifact_key: str, codes: torch.Tensor) -> None: + if _HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES <= 0 or _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES <= 0 or not artifact_key: + return + cached_codes = codes.detach().to("cpu", dtype=torch.long).contiguous() + size = int(cached_codes.numel() * cached_codes.element_size()) + if size > _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES: + return + previous = self._higgs_audio_v3_ref_code_cache.pop(artifact_key, None) + if previous is not None: + self._higgs_audio_v3_ref_code_cache_bytes -= previous[1] + self._higgs_audio_v3_ref_code_cache[artifact_key] = (cached_codes, size) + self._higgs_audio_v3_ref_code_cache_bytes += size + while len(self._higgs_audio_v3_ref_code_cache) > _HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES: + _, (_, old_size) = self._higgs_audio_v3_ref_code_cache.popitem(last=False) + self._higgs_audio_v3_ref_code_cache_bytes -= old_size + while self._higgs_audio_v3_ref_code_cache_bytes > _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES: + _, (_, old_size) = self._higgs_audio_v3_ref_code_cache.popitem(last=False) + self._higgs_audio_v3_ref_code_cache_bytes -= old_size + + async def _resolve_higgs_audio_v3_adapter(self): + """Lazy-load the tokenizer adapter for higgs_audio_v3.""" + cached = getattr(self, "_higgs_audio_v3_adapter", None) + if cached is not None: + return cached + + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + HiggsAudioV3TokenizerAdapter, + ) + + model_path = None + for stage in self.engine_client.stage_configs: + model_path = getattr(getattr(stage, "engine_args", None), "model", None) + if model_path: + break + if model_path is None: + model_path = getattr(self.engine_client, "model", None) + if model_path is None: + raise RuntimeError("higgs_audio_v3 serving could not resolve model path") + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + adapter = HiggsAudioV3TokenizerAdapter(tokenizer) + self._higgs_audio_v3_adapter = adapter + return adapter + def _validate_fish_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate Fish Speech request parameters. Returns error message or None.""" if not request.input or not request.input.strip(): @@ -1940,6 +2182,99 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None + def _validate_ming_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate Ming TTS request parameters. Returns error message or None.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if isinstance(request.ref_audio, list): + return self._validate_ming_tts_podcast_request(request) + return self._validate_ming_tts_single_speaker_request(request) + + def _validate_ming_tts_single_speaker_request(self, request: OpenAICreateSpeechRequest) -> str | None: + if request.ref_audio is not None: + fmt_err = self._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + + if request.speaker_embedding is not None: + if not request.speaker_embedding: + return "'speaker_embedding' must be a non-empty list of floats" + emb_len = len(request.speaker_embedding) + if emb_len != 192: + logger.warning( + "speaker_embedding has %d dimensions; Ming dense expects 192. " + "Wrong dimensions will likely fail or degrade output.", + emb_len, + ) + + voice_lower = request.voice.lower() if isinstance(request.voice, str) else None + uploaded_voice = bool(voice_lower and voice_lower in self.uploaded_speakers) + clone_source_present = request.ref_audio is not None or request.speaker_embedding is not None or uploaded_voice + + if request.task_type == "Base" and not clone_source_present: + return "Base task requires 'ref_audio', 'speaker_embedding', or an uploaded voice sample" + + if request.ref_audio is not None and request.ref_text is not None and not request.ref_text.strip(): + return "'ref_text' must be non-empty when provided with 'ref_audio'" + + # Ming offline ref-audio cases use prompt_waveform without prompt_text; + # keep the transcript requirement for other TTS models. + if request.ref_audio is not None and request.speaker_embedding is None and not self._is_ming_tts_model(): + uploaded_ref_text = self.uploaded_speakers[voice_lower].get("ref_text") if uploaded_voice else None + if not (request.ref_text and request.ref_text.strip()) and not uploaded_ref_text: + return "Reference-audio cloning requires non-empty 'ref_text'" + + if request.ref_text is not None and request.ref_audio is None and not uploaded_voice: + return "'ref_text' requires 'ref_audio' or an uploaded voice sample" + + if request.instructions and len(request.instructions) > self._max_instructions_length: + return f"Instructions too long (max {self._max_instructions_length} characters)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: + return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" + if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX: + return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}" + + return None + + def _validate_ming_tts_podcast_request(self, request: OpenAICreateSpeechRequest) -> str | None: + if len(request.ref_audio) < 2: + return "Podcast-style Ming requests require at least two 'ref_audio' clips" + + for ref_audio in request.ref_audio: + fmt_err = self._validate_ref_audio_format(ref_audio) + if fmt_err: + return fmt_err + + if not request.ref_text or not request.ref_text.strip(): + return "Podcast-style Ming requests require non-empty 'ref_text'" + + if request.speaker_embedding is not None: + embeddings = request.speaker_embedding + embedding_count = len(embeddings) if embeddings and isinstance(embeddings[0], list) else 1 + if embedding_count != len(request.ref_audio): + return ( + "Podcast-style Ming requests require one speaker embedding per ref_audio clip; " + f"got {embedding_count} embeddings for {len(request.ref_audio)} clips" + ) + if embeddings and isinstance(embeddings[0], list): + for item in embeddings: + if len(item) != 192: + return "Podcast-style Ming speaker embeddings must each have 192 dimensions" + + if request.instructions and len(request.instructions) > self._max_instructions_length: + return f"Instructions too long (max {self._max_instructions_length} characters)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: + return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" + if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX: + return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}" + + return None + async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple[list[float], int]: """Resolve ref_audio to (wav_samples, sample_rate). @@ -2066,6 +2401,142 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _discard_ref_audio_artifact_warmup(self, request_id: str) -> None: self._request_ref_audio_artifact_keys.pop(request_id, None) + async def _resolve_ref_audio_many(self, ref_audio_list: list[str]) -> list[tuple[list[float], int]]: + resolved = [] + for ref_audio in ref_audio_list: + resolved.append(await self._resolve_ref_audio(ref_audio)) + return resolved + + # ---- Ming TTS helpers ---- + + def _is_ming_tts_model(self) -> bool: + return self._tts_model_type == "ming_tts" + + def _coerce_ming_prompt_waveform(self, wav_samples, sample_rate): + from torchaudio.functional import resample as resample_audio + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import SAMPLE_RATE + + waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1) + if int(sample_rate) != SAMPLE_RATE: + waveform = resample_audio(waveform, int(sample_rate), SAMPLE_RATE) + return waveform + + def _build_ming_prompt_waveform( + self, + ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None, + ): + if isinstance(ref_audio_data, list): + return torch.cat( + [self._coerce_ming_prompt_waveform(item[0], item[1]) for item in ref_audio_data], + dim=-1, + ) + if ref_audio_data is not None: + return self._coerce_ming_prompt_waveform(ref_audio_data[0], ref_audio_data[1]) + return None + + def _extract_ming_speaker_embeddings_from_ref_audio( + self, + ref_audio_data_list: list[tuple[list[float], int]], + ) -> list[list[float]]: + from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor + + extractor = MingSpeakerEmbeddingExtractor(self.engine_client.model_config.model, target_sr=16000) + embeddings = [] + for wav_samples, sr in ref_audio_data_list: + waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1) + embedding = extractor.extract_from_waveform(waveform, int(sr)) + flat = embedding.detach().reshape(-1).to(torch.float32).cpu() + if int(flat.numel()) != 192: + raise ValueError(f"Ming speaker extractor returned {int(flat.numel())} dims; expected 192") + embeddings.append(flat.tolist()) + return embeddings + + def _parse_ming_instruction_fields( + self, + request, + *, + include_language=False, + include_voice=False, + plain_text_passthrough=False, + ): + instruction_text = request.instructions.strip() if isinstance(request.instructions, str) else None + instruction_dict: dict[str, Any] = {} + + if include_language and request.language not in (None, "", "Auto"): + instruction_dict["方言"] = request.language + + voice_lower = request.voice.lower() if isinstance(request.voice, str) else None + if include_voice and request.voice and not (voice_lower and voice_lower in self.uploaded_speakers): + instruction_dict["IP"] = request.voice + + if instruction_text: + try: + parsed = json.loads(instruction_text) + except json.JSONDecodeError: + parsed = None + if isinstance(parsed, dict): + instruction_dict.update(parsed) + elif instruction_dict or not plain_text_passthrough: + instruction_dict["风格"] = instruction_text + else: + return instruction_text + + return instruction_dict or None + + def _parse_ming_instruction(self, request: OpenAICreateSpeechRequest) -> Any: + """Build a Ming instruction payload from OpenAI speech fields.""" + return self._parse_ming_instruction_fields( + request, + include_language=True, + include_voice=True, + plain_text_passthrough=True, + ) + + def _build_ming_dense_prompt( + self, + request: OpenAICreateSpeechRequest, + *, + ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None = None, + ) -> dict[str, Any]: + """Build a Ming dense prompt directly from the OpenAI speech request.""" + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import KEY_MAX_DECODE_STEPS + from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt + + if self._tts_tokenizer is None: + model_name = self.engine_client.model_config.model + self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False) + + ref_text = request.ref_text + prompt_waveform = self._build_ming_prompt_waveform(ref_audio_data) if ref_text is not None else None + speaker_embedding = request.speaker_embedding + use_zero_spk_emb = prompt_waveform is None and speaker_embedding is None + + runtime_controls = {} + if request.max_new_tokens is not None: + runtime_controls[KEY_MAX_DECODE_STEPS] = request.max_new_tokens + + prompt_dict = build_ming_dense_prompt( + self._tts_tokenizer, + # bgm / music-prompt mode not supported online; + # requires prompt_mode API extension (deferred). + prompt=_MING_DEFAULT_PROMPT, + text=request.input, + runtime_controls=runtime_controls or None, + instruction=self._parse_ming_instruction(request), + prompt_text=ref_text, + prompt_waveform=prompt_waveform, + speaker_embedding=speaker_embedding, + use_zero_spk_emb=use_zero_spk_emb, + ) + prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"]) + prompt["prompt"] = prompt_dict["prompt"] + prompt["text"] = prompt_dict["text"] + prompt["additional_information"] = prompt_dict["additional_information"] + return prompt + async def _generate_audio_chunks( self, generator, @@ -2344,9 +2815,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if request.initial_codec_chunk_frames is not None: params["initial_codec_chunk_frames"] = [request.initial_codec_chunk_frames] - # VoiceDesign requires non_streaming_mode (match offline script behaviour). - # CustomVoice and Base rely on the model default (True and False respectively). - if params["task_type"][0] == "VoiceDesign": + if request.non_streaming_mode is not None: + params["non_streaming_mode"] = [request.non_streaming_mode] + # Preserve the legacy VoiceDesign fallback when the request omits an + # explicit override. CustomVoice and Base rely on model defaults + # (True and False respectively). + elif params["task_type"][0] == "VoiceDesign": params["non_streaming_mode"] = [True] return params @@ -2537,16 +3011,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): sampling_params_list = copy.deepcopy(sampling_params_list) hf_cfg = self.model_config.hf_config - model_path = self.engine_client.model_config.model - if not os.path.isdir(model_path): - from huggingface_hub import snapshot_download + # Build the Qwen tokenizer once per process (resolving the model dir via + # snapshot_download at most once) and reuse it across requests. + tokenizer = self._cosyvoice3_tokenizer + if tokenizer is None: + model_path = self.engine_client.model_config.model + if not os.path.isdir(model_path): + from huggingface_hub import snapshot_download - model_path = snapshot_download(model_path) - tokenizer = get_qwen_tokenizer( - token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path), - skip_special_tokens=hf_cfg.skip_special_tokens, - version=hf_cfg.version, - ) + model_path = snapshot_download(model_path) + tokenizer = get_qwen_tokenizer( + token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path), + skip_special_tokens=hf_cfg.skip_special_tokens, + version=hf_cfg.version, + ) + self._cosyvoice3_tokenizer = tokenizer _, text_token_len = extract_text_token( request.input, tokenizer, @@ -2665,25 +3144,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # ---- Ming-flash-omni standalone-talker (TTS) helpers ---- - def _build_ming_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: + def _build_ming_flash_omni_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: # request.instructions accepts two forms: # 1. Plain text: mapped to the caption's 风格 (style) field # 2. JSON object: parsed and splatted into the caption. Unlocks # Unknown keys are dropped by `ming_create_instruction`. - caption_fields: dict[str, Any] = {} - if request.instructions: - stripped = request.instructions.strip() - if stripped.startswith("{"): - try: - parsed = json.loads(stripped) - except json.JSONDecodeError: - parsed = None - if isinstance(parsed, dict): - caption_fields.update(parsed) - else: - caption_fields["风格"] = request.instructions - else: - caption_fields["风格"] = request.instructions + caption_fields = self._parse_ming_instruction_fields(request) or {} has_spk_emb = request.speaker_embedding is not None @@ -2748,6 +3214,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request, ref_audio_data=ref_audio_data, has_inline_ref_audio=has_inline_ref_audio ) tts_params = {} + # Fish Speech resolves conditioning entirely from raw request fields + # (text / ref_audio / ref_text rebuilt into inputs_embeds); unlike + # uploaded-voice models it has no resolved tts_params to fold in, so + # the salt is derived from the request alone. + prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params) elif self._tts_model_type == "omnivoice": if not request.input or not request.input.strip(): raise ValueError("Input text cannot be empty") @@ -2796,6 +3267,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): additional["voice_name"] = voice_lower additional["voice_created_at"] = self._voice_created_at(voice_lower) tts_params = {} + elif self._tts_model_type == "higgs_audio_v3": + err = self._apply_uploaded_speaker(request) + if err: + raise ValueError(err) + validation_error = self._validate_higgs_audio_v3_request(request) + if validation_error: + raise ValueError(validation_error) + prompt = await self._build_higgs_audio_v3_params(request) + tts_params = {} elif self._tts_model_type == "voxcpm2": # voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the # raw waveform tuple for prefill-length accounting, not a base64 data URL. @@ -2837,8 +3317,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prompt = await self._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio) tts_params = {} elif self._tts_model_type == "ming_flash_omni_tts": - prompt = self._build_ming_prompt(request) + prompt = self._build_ming_flash_omni_prompt(request) tts_params = {} + elif self._tts_model_type == "ming_tts": + ref_audio_source = request.ref_audio + voice_lower = request.voice.lower() if isinstance(request.voice, str) else None + if ref_audio_source is None and voice_lower in self.uploaded_speakers: + ref_audio_source = self._get_uploaded_audio_data(request.voice) + if request.ref_text is None: + request.ref_text = self.uploaded_speakers[voice_lower].get("ref_text") + ref_audio_data = None + if isinstance(ref_audio_source, list): + ref_audio_data = await self._resolve_ref_audio_many(ref_audio_source) + if request.speaker_embedding is None: + request.speaker_embedding = self._extract_ming_speaker_embeddings_from_ref_audio(ref_audio_data) + elif ref_audio_source is not None and isinstance(ref_audio_source, str): + wav_list, sr = await self._resolve_ref_audio(ref_audio_source) + ref_audio_data = (wav_list, sr) + if request.speaker_embedding is None: + request.speaker_embedding = self._extract_ming_speaker_embeddings_from_ref_audio( + [ref_audio_data] + )[0] + prompt = self._build_ming_dense_prompt(request, ref_audio_data=ref_audio_data) + tts_params = prompt.get("additional_information", {}) elif self._tts_model_type in ("moss_tts_nano", "moss_tts"): tts_params = await self._build_moss_tts_params(request) if request.voice: @@ -2924,6 +3425,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "voxcpm2" elif self._tts_model_type == "ming_flash_omni_tts": model_type = "ming_flash_omni_tts" + elif self._tts_model_type == "ming_tts": + model_type = "ming_tts" elif self._tts_model_type == "moss_tts_nano": model_type = "moss_tts_nano" elif self._tts_model_type == "moss_tts": @@ -3032,6 +3535,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): getattr(sampling_params_list[0], "min_tokens", 0), request.max_new_tokens, ) + elif self._tts_model_type == "ming_tts" and sampling_params_list: + import copy + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import TEXT_EOS_TOKEN_ID + + sampling_params_list = copy.deepcopy(sampling_params_list) + sampling_params_list[0].stop_token_ids = [int(TEXT_EOS_TOKEN_ID)] + if request.max_new_tokens is not None: + # Ming emits TEXT_EOS after the latent decode budget is exhausted, so + # Stage-0 needs one extra token beyond ming_max_decode_steps. + sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1 if request.seed is not None and sampling_params_list: if sampling_params_list is self.engine_client.default_sampling_params_list: @@ -3463,6 +3977,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): x_vector_only_mode=_pick("x_vector_only_mode"), max_new_tokens=_pick("max_new_tokens"), initial_codec_chunk_frames=_pick("initial_codec_chunk_frames"), + non_streaming_mode=_pick("non_streaming_mode"), ) async def create_speech_batch( diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 695df04c2..695e87d6e 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -221,6 +221,7 @@ class OmniStreamingSpeechHandler: speed=config.speed, max_new_tokens=config.max_new_tokens, initial_codec_chunk_frames=config.initial_codec_chunk_frames, + non_streaming_mode=config.non_streaming_mode, ref_audio=config.ref_audio, ref_text=config.ref_text, x_vector_only_mode=config.x_vector_only_mode, diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 8673e1923..f05ce204b 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -39,6 +39,20 @@ class ReferenceImage: data: Image.Image +@dataclass +class ReferenceVideo: + """Reference video frames for video-conditioned generation.""" + + data: list[Image.Image] + + +@dataclass +class ReferenceAudio: + """Reference audio file path for speech-to-video generation.""" + + path: str + + @dataclass class VideoGenerationArtifacts: """Normalized outputs and profiler metadata extracted from one request.""" @@ -96,6 +110,8 @@ class OmniOpenAIServingVideo: reference_id: str, *, reference_image: ReferenceImage | None = None, + reference_video: ReferenceVideo | None = None, + reference_audio: ReferenceAudio | None = None, ) -> VideoGenerationArtifacts: """Run the generation pipeline and extract video/audio/profiler outputs.""" prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt, modalities=["video"]) @@ -105,13 +121,26 @@ class OmniOpenAIServingVideo: gen_params = self._resolve_default_sampling_params() input_image = None if reference_image is None else reference_image.data + input_video = None if reference_video is None else reference_video.data + if input_image is not None and input_video is not None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Provide either an image reference or a video reference, not both.", + ) vp = request.resolve_video_params() if input_image is not None and vp.width is not None and vp.height is not None: target_size = (vp.width, vp.height) if input_image.size != target_size: input_image = input_image.resize(target_size, Image.Resampling.LANCZOS) + multi_modal_data: dict[str, Any] = {} if input_image is not None: - prompt["multi_modal_data"] = {"image": input_image} + multi_modal_data["image"] = input_image + elif input_video is not None: + multi_modal_data["video"] = input_video + if reference_audio is not None: + multi_modal_data["audio"] = reference_audio.path + if multi_modal_data: + prompt["multi_modal_data"] = multi_modal_data if vp.width is not None and vp.height is not None: gen_params.width = vp.width gen_params.height = vp.height @@ -211,8 +240,16 @@ class OmniOpenAIServingVideo: reference_id: str, *, reference_image: ReferenceImage | None = None, + reference_video: ReferenceVideo | None = None, + reference_audio: ReferenceAudio | None = None, ) -> VideoGenerationResponse: - artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image) + artifacts = await self._run_and_extract( + request, + reference_id, + reference_image=reference_image, + reference_video=reference_video, + reference_audio=reference_audio, + ) video_codec_options = {"preset": "ultrafast", "threads": "0"} if request.extra_params is not None and isinstance(request.extra_params, dict): @@ -252,9 +289,17 @@ class OmniOpenAIServingVideo: reference_id: str, *, reference_image: ReferenceImage | None = None, + reference_video: ReferenceVideo | None = None, + reference_audio: ReferenceAudio | None = None, ) -> tuple[bytes, dict[str, float], float, VideoAction | None]: """Generate a video and return raw MP4 bytes, bypassing base64 encoding.""" - artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image) + artifacts = await self._run_and_extract( + request, + reference_id, + reference_image=reference_image, + reference_video=reference_video, + reference_audio=reference_audio, + ) if len(artifacts.videos) > 1: logger.warning( "Video request %s generated %d outputs; returning only the first.", diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 303b3b4a1..aa4eeaba5 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -8,8 +8,9 @@ from __future__ import annotations import base64 import binascii +from collections import deque from io import BytesIO -from typing import Any +from typing import Any, Literal import httpx import numpy as np @@ -19,8 +20,11 @@ from PIL import Image, UnidentifiedImageError from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError from vllm_omni.entrypoints.openai.protocol.videos import ( FileImageReference, + FileVideoReference, ImageReference, UrlImageReference, + UrlVideoReference, + VideoReference, ) @@ -31,6 +35,70 @@ def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image: raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc +def _decode_video_bytes( + video_bytes: bytes, + *, + source: str, + max_frames: int | None = None, + keep: Literal["first", "last"] = "first", +) -> list[Image.Image]: + try: + import av + except ImportError as exc: # pragma: no cover - av is a serving dependency via media_utils + raise InvalidInputReferenceError(f"Invalid {source}: video decoding requires PyAV.") from exc + + if keep not in {"first", "last"}: + raise InvalidInputReferenceError(f"Invalid {source}: video frame selection must be 'first' or 'last'.") + if max_frames is not None and max_frames <= 0: + raise InvalidInputReferenceError(f"Invalid {source}: max video frames must be positive.") + + frames: list[Image.Image] = [] + tail_frames: deque[Image.Image] | None = ( + deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None + ) + try: + with av.open(BytesIO(video_bytes)) as container: + for frame in container.decode(video=0): + image = frame.to_image().convert("RGB") + if tail_frames is not None: + tail_frames.append(image) + else: + frames.append(image) + if keep == "first" and max_frames is not None and len(frames) >= max_frames: + break + except Exception as exc: + raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") from exc + + if tail_frames is not None: + frames = list(tail_frames) + if not frames: + raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") + return frames + + +def _decode_media_bytes( + media_bytes: bytes, + *, + source: str, + max_video_frames: int | None = None, + video_keep: Literal["first", "last"] = "first", +) -> Image.Image | list[Image.Image]: + try: + return _decode_image_bytes(media_bytes, source=source) + except InvalidInputReferenceError: + try: + return _decode_video_bytes( + media_bytes, + source=source, + max_frames=max_video_frames, + keep=video_keep, + ) + except InvalidInputReferenceError as video_exc: + raise InvalidInputReferenceError( + f"Invalid {source}: provided content is not a valid image or video." + ) from video_exc + + def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image: if input_reference: if input_reference.startswith("data:image"): @@ -64,23 +132,143 @@ async def decode_image_url(image_url: str) -> Image.Image: raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.") +def _decode_base64_video( + video_reference: str, + *, + source: str, + max_frames: int | None = None, + keep: Literal["first", "last"] = "first", +) -> list[Image.Image]: + if video_reference: + if video_reference.startswith("data:video"): + _, b64_data = video_reference.split(",", 1) + else: + b64_data = video_reference + + try: + video_bytes = base64.b64decode(b64_data) + except (binascii.Error, ValueError) as exc: # pragma: no cover - malformed base64 + raise InvalidInputReferenceError(f"Invalid {source}: video data is not valid base64.") from exc + return _decode_video_bytes(video_bytes, source=source, max_frames=max_frames, keep=keep) + raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.") + + +async def decode_video_url( + video_url: str, + *, + max_frames: int | None = None, + keep: Literal["first", "last"] = "first", +) -> list[Image.Image]: + if video_url.startswith("data:video"): + return _decode_base64_video( + video_url, + source="video_reference.video_url", + max_frames=max_frames, + keep=keep, + ) + + if video_url.startswith(("http://", "https://")): + async with httpx.AsyncClient(timeout=60) as client: + try: + response = await client.get(video_url) + response.raise_for_status() + except httpx.HTTPError as exc: + raise InvalidInputReferenceError( + "Invalid video_reference.video_url: failed to download video." + ) from exc + return _decode_video_bytes( + response.content, + source="video_reference.video_url", + max_frames=max_frames, + keep=keep, + ) + + raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.") + + +async def decode_audio_url(audio_url: str) -> str: + """Decode an audio URL or data-URL to a temporary file path.""" + import tempfile + + audio_bytes: bytes | None = None + + if audio_url.startswith("data:audio"): + _, b64_data = audio_url.split(",", 1) + try: + audio_bytes = base64.b64decode(b64_data) + except (binascii.Error, ValueError) as exc: + raise InvalidInputReferenceError( + "Invalid audio_reference.audio_url: audio data is not valid base64." + ) from exc + elif audio_url.startswith(("http://", "https://")): + async with httpx.AsyncClient(timeout=60) as client: + try: + response = await client.get(audio_url) + response.raise_for_status() + except httpx.HTTPError as exc: + raise InvalidInputReferenceError( + "Invalid audio_reference.audio_url: failed to download audio." + ) from exc + audio_bytes = response.content + else: + raise InvalidInputReferenceError("Invalid audio_reference.audio_url: must be an http(s) URL or data URL.") + + if not audio_bytes: + raise InvalidInputReferenceError("Invalid audio_reference: audio data is empty.") + + suffix = ".wav" + if audio_url.startswith("data:audio/"): + mime = audio_url.split(";")[0].removeprefix("data:") + ext = mime.split("/")[-1] + if ext in ("mpeg", "mp3"): + suffix = ".mp3" + elif ext == "wav": + suffix = ".wav" + elif ext.isalnum() and len(ext) <= 8: + suffix = f".{ext}" + + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) + tmp.write(audio_bytes) + tmp.close() + return tmp.name + + async def decode_input_reference( image_reference: ImageReference | None, + video_reference: VideoReference | None, input_reference_bytes: bytes | None, -) -> Image.Image | None: - """Decode image input from multipart bytes, base64/data URL, or image_reference.""" + *, + max_video_frames: int | None = None, + video_keep: Literal["first", "last"] = "first", +) -> Image.Image | list[Image.Image] | None: + """Decode media input from multipart bytes, data URLs, or typed references.""" - if input_reference_bytes is not None and image_reference is not None: - raise InvalidInputReferenceError("Provide either input_reference or image_reference, not both.") + provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference)) + if provided > 1: + raise InvalidInputReferenceError("Provide only one of input_reference, image_reference, or video_reference.") if isinstance(input_reference_bytes, bytes): - return _decode_image_bytes(input_reference_bytes, source="input_reference") + return _decode_media_bytes( + input_reference_bytes, + source="input_reference", + max_video_frames=max_video_frames, + video_keep=video_keep, + ) if isinstance(image_reference, UrlImageReference): return await decode_image_url(image_reference.image_url) elif isinstance(image_reference, FileImageReference): raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.") + if isinstance(video_reference, UrlVideoReference): + return await decode_video_url( + video_reference.video_url, + max_frames=max_video_frames, + keep=video_keep, + ) + elif isinstance(video_reference, FileVideoReference): + raise InvalidInputReferenceError("Invalid video_reference: file_id is not supported yet.") + return None diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index e7c91cf4d..74b533b51 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -95,24 +95,6 @@ def _map_device_list(stage_id: int, device_list: list[str], visible_device_list: if not all(device.isdigit() for device in device_list): raise ValueError("Logical devices must be non-negative integers") - # Idempotency: if every requested id already names an entry in the visible - # pool (by value, not by index), treat the YAML field as physical and skip - # the remap. This avoids a double-mapping crash when a parent harness has - # already narrowed CUDA_VISIBLE_DEVICES to specific physical ids before - # spawning the subprocess (the test runtime's OmniServerStageCli does this - # — it sets each stage subprocess's env via ``_set_stage_device_env``). - # Without this guard, a YAML ``devices: "1"`` against a narrowed visible - # list ``["1"]`` would try ``visible[1]`` (out of range) and raise. - visible_set = set(visible_device_list) - if all(device in visible_set for device in device_list): - logger.info( - "Stage %s requested devices %s are already in the visible set %s; using as-is (no remap).", - stage_id, - device_list, - visible_device_list, - ) - return list(device_list) - logical_ids = [int(device) for device in device_list] # If ALL logical device IDs are >= num_visible, none can be 0-based indices diff --git a/vllm_omni/errors.py b/vllm_omni/errors.py new file mode 100644 index 000000000..1bc7f5913 --- /dev/null +++ b/vllm_omni/errors.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Request-scoped client error types shared across vLLM-Omni entrypoints.""" + +from __future__ import annotations + +from collections.abc import Callable +from http import HTTPStatus +from typing import NoReturn + +DEFAULT_CLIENT_ERROR_TYPE = "BadRequestError" + + +class OmniClientError(ValueError): + """ + Request-scoped error that should be surfaced as a 4xx response. + One example of using OmniClientError is GuardrailViolationError, which is captured and resurfaced + as HTTP 400 error code, instead of a generic 500. OmniClientError should be used for any exceptions + which need to be resurfaced as 4xx, as opposed to EngineDeadError/EngineGenerateError which are resurfaced as 500. + """ + + def __init__( + self, + message: str, + *, + status_code: int = HTTPStatus.BAD_REQUEST.value, + error_type: str = DEFAULT_CLIENT_ERROR_TYPE, + ) -> None: + super().__init__(message) + self.message = message + self.status_code = int(status_code) + self.error_type = error_type + + +class GuardrailViolationError(OmniClientError): + """Raised when a model guardrail rejects request content.""" + + +def client_error_metadata(exc: BaseException) -> tuple[int | None, str | None]: + if isinstance(exc, OmniClientError): + return exc.status_code, exc.error_type + return None, None + + +def client_error_from_metadata( + message: str, + *, + status_code: int | None, + error_type: str | None, +) -> OmniClientError: + return OmniClientError( + message, + status_code=status_code or HTTPStatus.BAD_REQUEST.value, + error_type=error_type or DEFAULT_CLIENT_ERROR_TYPE, + ) + + +def is_client_error_status(status_code: int | None) -> bool: + return status_code is not None and 400 <= int(status_code) < 500 + + +def raise_client_error_or( + message: str, + *, + status_code: int | None, + error_type: str | None, + fallback: Callable[[str], BaseException], +) -> NoReturn: + """Raise a client error for 4xx statuses, otherwise raise ``fallback(message)``. + + Centralizes the "client-error-or-fallback" decision shared by the engine + error paths so the status mapping lives in one place. + """ + if is_client_error_status(status_code): + raise client_error_from_metadata( + message, + status_code=status_code, + error_type=error_type, + ) + raise fallback(message) diff --git a/vllm_omni/metrics/utils.py b/vllm_omni/metrics/utils.py index 1ed987a19..2062eb5ba 100644 --- a/vllm_omni/metrics/utils.py +++ b/vllm_omni/metrics/utils.py @@ -94,7 +94,13 @@ def _format_table( if isinstance(value, float): return f"{value:,.3f}" if isinstance(value, list): - return ", ".join(str(f"{v:,.3f}") for v in value) + if not value: + return "" + try: + avg = sum(float(v) for v in value) / len(value) + return f"{avg:,.3f} (n={len(value)})" + except (TypeError, ValueError): + return ", ".join(str(v) for v in value) return str(value) table = PrettyTable() @@ -117,6 +123,17 @@ def _format_table( if column_key is None: raise ValueError("column_key is required for multi-column mode") col_headers = [f"{column_prefix}{row.get(column_key, '?')}" for row in data] + # PrettyTable requires unique field names. When the same column key + # appears more than once (e.g. two events with the same stage_id), + # deduplicate by appending _2, _3, … to the later occurrences. + if len(col_headers) != len(set(col_headers)): + seen: dict[str, int] = {} + deduped: list[str] = [] + for h in col_headers: + n = seen.get(h, 0) + seen[h] = n + 1 + deduped.append(h if n == 0 else f"{h}_{n + 1}") + col_headers = deduped table.field_names = ["Field"] + col_headers table.align["Field"] = "l" for col in col_headers: diff --git a/vllm_omni/model_executor/models/common/ming/__init__.py b/vllm_omni/model_executor/models/common/ming/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/common/ming/audio_dsp.py b/vllm_omni/model_executor/models/common/ming/audio_dsp.py new file mode 100644 index 000000000..be8cd489c --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/audio_dsp.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/istft.py + +import torch +import torch.nn as nn + + +class ISTFT(nn.Module): + """ + Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with + windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges. + See issue: https://github.com/pytorch/pytorch/issues/62323 + Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs. + The NOLA constraint is met as we trim padded samples anyway. + + Args: + n_fft (int): Size of Fourier transform. + hop_length (int): The distance between neighboring sliding window frames. + win_length (int): The size of window frame and STFT filter. + padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". + """ + + def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): + super().__init__() + if padding not in ["center", "same"]: + raise ValueError("Padding must be 'center' or 'same'.") + self.padding = padding + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + window = torch.hann_window(win_length) + self.register_buffer("window", window) + + self.buffer_len = self.win_length - self.hop_length + + def _buffer_process( + self, + x: torch.Tensor, + buffer: torch.Tensor | None, + pad: int, + last_chunk: bool = False, + streaming: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if streaming: + if buffer is None: + # first chunk + x = x[:, pad:] + if buffer is not None: + # next chunk + x[:, : self.buffer_len] += buffer + buffer = x[:, -self.buffer_len :] + if not last_chunk: + x = x[:, : -self.buffer_len] + else: + x = x[:, :-pad] + else: + x = x[:, pad:-pad] + + return x, buffer + + def forward( + self, + spec: torch.Tensor, + audio_buffer: torch.Tensor | None = None, + window_buffer: torch.Tensor | None = None, + streaming: bool = False, + last_chunk: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """ + Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram. + + Args: + spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size, + N is the number of frequency bins, and T is the number of time frames. + audio_buffer (Tensor): [Streaming Input/State] The audio overlap buffer from the previous chunk. + Shape: (B, win_length - hop_length) + window_buffer (Tensor): [Streaming Input/State] The window overlap buffer from the previous chunk. + streaming: If `True`, the function operates in streaming mode, processing `spec` as a single chunk. + last_chunk: When `streaming=True` and `last_chunk=True`, the function can perform final "flush" operations + + Returns: + Reconstructed signal, plus streaming buffers when `padding="same"`. + """ + if self.padding == "center": + # Fallback to pytorch native implementation + return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True) + elif self.padding == "same": + pad = (self.win_length - self.hop_length) // 2 + else: + raise ValueError("Padding must be 'center' or 'same'.") + + if spec.dim() != 3: + raise ValueError(f"Expected spec rank-3 [Batch, Freq, Time], got {tuple(spec.shape)}") + _, _, T = spec.shape + + # Inverse FFT + ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") + ifft = ifft * self.window[None, :, None] + + # Overlap and Add + output_size = (T - 1) * self.hop_length + self.win_length + y = torch.nn.functional.fold( + ifft, + output_size=(1, output_size), + kernel_size=(1, self.win_length), + stride=(1, self.hop_length), + )[:, 0, 0, :] + + y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming) + + # Window envelope + window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) + window_envelope = ( + torch.nn.functional.fold( + window_sq, + output_size=(1, output_size), + kernel_size=(1, self.win_length), + stride=(1, self.hop_length), + ) + .squeeze(0) + .squeeze(0) + ) + + window_envelope, window_buffer = self._buffer_process( + window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming + ) + window_envelope = window_envelope.squeeze() + + # Normalize + if not (window_envelope > 1e-11).all(): + raise RuntimeError("ISTFT window envelope underflowed; invalid overlap-add state.") + y = y / window_envelope + + return y, audio_buffer, window_buffer + + +class FourierHead(nn.Module): + """Base class for inverse fourier modules.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, + L is the sequence length, and H denotes the model dimension. + + Returns: + Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. + """ + raise NotImplementedError("Subclasses must implement the forward method.") + + +class ISTFTHead(FourierHead): + """ + ISTFT Head module for predicting STFT complex coefficients. + + Args: + dim (int): Hidden dimension of the model. + n_fft (int): Size of Fourier transform. + hop_length (int): The distance between neighboring sliding window frames, which should align with + the resolution of the input features. + padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". + """ + + def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): + super().__init__() + out_dim = n_fft + 2 + self.out = torch.nn.Linear(dim, out_dim) + self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding) + + def forward( + self, + x: torch.Tensor, + audio_buffer: torch.Tensor | None = None, + window_buffer: torch.Tensor | None = None, + streaming: bool = False, + last_chunk: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """ + Forward pass of the ISTFTHead module. + + Args: + x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, + L is the sequence length, and H denotes the model dimension. + + Returns: + Audio, predicted spectrogram coefficients, and streaming buffers. + """ + x_pred = self.out(x) + x_pred = x_pred.transpose(1, 2) + mag, p = x_pred.chunk(2, dim=1) + mag = torch.exp(mag) + mag = torch.clip(mag, max=1e2) # safeguard to prevent excessively large magnitudes + # wrapping happens here. These two lines produce real and imaginary value + x = torch.cos(p) + y = torch.sin(p) + # recalculating phase here does not produce anything new + # only costs time + # phase = torch.atan2(y, x) + # S = mag * torch.exp(phase * 1j) + # better directly produce the complex value + S = mag * (x + 1j * y) + audio, audio_buffer, window_buffer = self.istft( + S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk + ) + return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer + + +__all__ = ["FourierHead", "ISTFT", "ISTFTHead"] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py b/vllm_omni/model_executor/models/common/ming/audio_vae.py similarity index 64% rename from vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py rename to vllm_omni/model_executor/models/common/ming/audio_vae.py index 3137a878e..6924ae548 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py +++ b/vllm_omni/model_executor/models/common/ming/audio_vae.py @@ -1,10 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright 2025 The vLLM-Omni team. -# Copyright (c) Ant Group. All rights reserved. -# Adapted from: -# https://github.com/inclusionAI/Ming/tree/e58533db227031990c5a6864dcf5f08fb53ed0d2/AudioVAE - -from __future__ import annotations +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch import torch.nn as nn @@ -13,6 +8,8 @@ from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Mo from transformers.utils import is_flash_attn_2_available from vllm.logger import init_logger +from .audio_dsp import ISTFTHead + logger = init_logger(__name__) @@ -21,114 +18,50 @@ class AudioVAEConfig(PretrainedConfig): def __init__( self, - sample_rate: int = 44100, - enc_kwargs: dict | None = None, - dec_kwargs: dict | None = None, - init_method: str = "kaiming", - patch_size: int = 4, + sample_rate=44100, + enc_kwargs=None, + semantic_module_kwargs=None, + dec_kwargs=None, + hifi_gan_disc_kwargs=None, + spec_disc_kwargs=None, + lambda_disc=1.0, + lambda_mel_loss=15, + lambda_adv=1.0, + lambda_feat_match_loss=1.0, + lambda_semantic=5.0, + init_method="kaiming", + patch_size=4, **kwargs, ): self.sample_rate = sample_rate self.enc_kwargs = enc_kwargs or {} + self.semantic_module_kwargs = semantic_module_kwargs self.dec_kwargs = dec_kwargs or {} + self.hifi_gan_disc_kwargs = hifi_gan_disc_kwargs + self.spec_disc_kwargs = spec_disc_kwargs + self.lambda_disc = lambda_disc + self.lambda_mel_loss = lambda_mel_loss + self.lambda_adv = lambda_adv + self.lambda_feat_match_loss = lambda_feat_match_loss + self.lambda_semantic = lambda_semantic self.init_method = init_method self.patch_size = patch_size super().__init__(**kwargs) -class ISTFT(nn.Module): - def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): - super().__init__() - if padding not in ["center", "same"]: - raise ValueError("Padding must be 'center' or 'same'.") - self.padding = padding - self.n_fft = n_fft - self.hop_length = hop_length - self.win_length = win_length - window = torch.hann_window(win_length) - self.register_buffer("window", window) - self.buffer_len = self.win_length - self.hop_length - - def _buffer_process(self, x, buffer, pad, last_chunk=False, streaming=False): - if streaming: - if buffer is None: - x = x[:, pad:] - if buffer is not None: - x[:, : self.buffer_len] += buffer - buffer = x[:, -self.buffer_len :] - if not last_chunk: - x = x[:, : -self.buffer_len] - else: - x = x[:, :-pad] - else: - x = x[:, pad:-pad] - return x, buffer - - def forward(self, spec, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False): - if self.padding == "center": - return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True) - elif self.padding == "same": - pad = (self.win_length - self.hop_length) // 2 - else: - raise ValueError("Padding must be 'center' or 'same'.") - - B, N, T = spec.shape - ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") - ifft = ifft * self.window[None, :, None] - - output_size = (T - 1) * self.hop_length + self.win_length - y = torch.nn.functional.fold( - ifft, - output_size=(1, output_size), - kernel_size=(1, self.win_length), - stride=(1, self.hop_length), - )[:, 0, 0, :] - - y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming) - - window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) - window_envelope = ( - torch.nn.functional.fold( - window_sq, - output_size=(1, output_size), - kernel_size=(1, self.win_length), - stride=(1, self.hop_length), - ) - .squeeze(0) - .squeeze(0) - ) - - window_envelope, window_buffer = self._buffer_process( - window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming - ) - window_envelope = window_envelope.squeeze() - - assert (window_envelope > 1e-11).all() - y = y / window_envelope - - return y, audio_buffer, window_buffer - - -class ISTFTHead(nn.Module): - def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): - super().__init__() - out_dim = n_fft + 2 - self.out = nn.Linear(dim, out_dim) - self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding) - - def forward(self, x, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False): - x_pred = self.out(x) - x_pred = x_pred.transpose(1, 2) - mag, p = x_pred.chunk(2, dim=1) - mag = torch.exp(mag) - mag = torch.clip(mag, max=1e2) - x = torch.cos(p) - y = torch.sin(p) - S = mag * (x + 1j * y) - audio, audio_buffer, window_buffer = self.istft( - S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk - ) - return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer +def _qwen2_config(backbone): + config = Qwen2Config.from_dict(config_dict=backbone) + if (getattr(config, "_attn_implementation", None) or getattr(config, "attn_implementation", None)) not in ( + None, + "flash_attention_2", + ): + return config + if is_flash_attn_2_available(): + config._attn_implementation_autoset = True + config._attn_implementation = "flash_attention_2" + else: + config._attn_implementation = "sdpa" + return config class StreamingLinearUpsample(nn.Module): @@ -138,6 +71,8 @@ class StreamingLinearUpsample(nn.Module): self.upsampler = nn.Upsample(scale_factor=scale_factor, mode="linear", align_corners=False) def forward(self, x, state=None, is_last=False): + if x is None and is_last and (state is None or state.get("prev_chunk") is None): + raise ValueError("Received end-of-stream without any latent chunk to upsample.") if state is None: state = {"prev_chunk": None, "history_last": None, "is_first": True} @@ -188,16 +123,72 @@ class StreamingLinearUpsample(nn.Module): return final_out, state +class Encoder(nn.Module): + def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1): + super().__init__() + config = _qwen2_config(encoder_args) + logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation) + self.encoder = Qwen2Model(config) + self.input_dim = input_dim + self.hop_size = hop_size + self.latent_dim = latent_dim + self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False) + self.fc2 = nn.Linear(config.hidden_size, config.hidden_size) + self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2) + self.norm = nn.LayerNorm(config.hidden_size) + self.patch_size = patch_size + if patch_size != -1: + aggregator_config = _qwen2_config({**encoder_args, "num_hidden_layers": 4}) + self.aggregator = Qwen2Model(aggregator_config) + self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size)) + self.cls_embed.data.normal_(0, 0.02) + + def get_frames(self, x): + num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size + expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim + padding_needed = expected_len - x.size(-1) + waveform = F.pad(x, (0, padding_needed), value=0.0) + frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size) + return frames + + def pad_patch_insert_cls(self, x): + bsz, _, dim = x.size() + num_frame = x.size(1) + r = num_frame % self.patch_size + pad_num = self.patch_size - r if r else 0 + x = F.pad(x, (0, 0, 0, pad_num), value=0.0) + # [Batch, Time, Dimension] -> [Batch*PatchGroups, Patch, Dimension]. + x = x.reshape(-1, self.patch_size, dim) + x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1) + # [Batch*PatchGroups, Patch+1, Dimension] -> [Batch, Time, Dimension]. + x = x.reshape(bsz, -1, dim) + return x + + def forward(self, waveform): + x = self.get_frames(waveform) + x = self.fc1(x) + x = self.fc2(x) + x = self.encoder(inputs_embeds=x) + x = x.last_hidden_state + + if self.patch_size != -1: + x = self.pad_patch_insert_cls(x) + x = self.aggregator(inputs_embeds=x) + x = x.last_hidden_state + bsz, _, dim = x.size() + # [Batch, Time, Dimension] -> [Batch*PatchGroups, Patch+1, Dimension]. + x = x.reshape(-1, self.patch_size + 1, dim) + # [Batch*PatchGroups, 1, Dimension] -> [Batch, PatchGroups, Dimension]. + x = x[:, -1:, :].reshape(bsz, -1, dim) + + x = self.fc3(x) + return x, waveform.unsqueeze(1) + + class Decoder(nn.Module): def __init__(self, decoder_args, output_dim=320, latent_dim=64, patch_size=-1): super().__init__() - config = Qwen2Config.from_dict(config_dict=decoder_args) - if is_flash_attn_2_available(): - config._attn_implementation_autoset = True - config._attn_implementation = "flash_attention_2" - else: - config._attn_implementation = "sdpa" - + config = _qwen2_config(decoder_args) logger.info("AudioVAE Decoder: using attn_implementation=%r", config._attn_implementation) self.decoder = Qwen2Model(config) self.output_dim = output_dim @@ -211,7 +202,21 @@ class Decoder(nn.Module): if self.patch_size != -1: self.upsampling = StreamingLinearUpsample(scale_factor=patch_size) + def forward(self, x): + x = self.fc1(x) + + if self.patch_size != -1: + # [Batch, Time, Dimension] -> [Batch, Dimension, Time] -> [Batch, Time, Dimension]. + x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2) + + x = self.decoder(inputs_embeds=x) + x = x.last_hidden_state + x, _ = self.head(x) + return x, None + def low_level_reconstruct(self, x, past_key_values=None, use_cache=False, stream_state=None, last_chunk=False): + if stream_state is None: + stream_state = (None, None, None) upsample_state, audio_buffer, window_buffer = stream_state bsz, device, dtype = x.size(0), x.device, x.dtype x = self.fc1(x) @@ -222,6 +227,7 @@ class Decoder(nn.Module): stream_state = (upsample_state, audio_buffer, window_buffer) return torch.empty(bsz, 1, 0, device=device, dtype=dtype), stream_state, past_key_values else: + # [Batch, Time, Dimension] -> [Batch, Dimension, Time] -> [Batch, Time, Dimension]. x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2) hidden_states_list = [] @@ -239,7 +245,6 @@ class Decoder(nn.Module): past_len = 0 curr_len = x.shape[1] - if past_len < target_len and (past_len + curr_len) >= sw_size: fill_len = target_len - past_len x_fill = x[:, :fill_len, :] @@ -269,86 +274,39 @@ class Decoder(nn.Module): return x_out, stream_state, past_key_values -class Encoder(nn.Module): - def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1): - super().__init__() - config = Qwen2Config.from_dict(config_dict=encoder_args) - if is_flash_attn_2_available(): - config._attn_implementation_autoset = True - config._attn_implementation = "flash_attention_2" - else: - config._attn_implementation = "sdpa" - - logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation) - self.encoder = Qwen2Model(config) - self.input_dim = input_dim - self.hop_size = hop_size - self.latent_dim = latent_dim - self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False) - self.fc2 = nn.Linear(config.hidden_size, config.hidden_size) - self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2) - self.norm = nn.LayerNorm(config.hidden_size) - self.patch_size = patch_size - if patch_size != -1: - config.num_hidden_layers = 4 - self.aggregator = Qwen2Model(config) - self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size)) - self.cls_embed.data.normal_(0, 0.02) - - def get_frames(self, x): - num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size - expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim - padding_needed = expected_len - x.size(-1) - waveform = F.pad(x, (0, padding_needed), value=0.0) - frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size) - return frames - - def pad_patch_insert_cls(self, x): - bsz, _, dim = x.size() - num_frame = x.size(1) - r = num_frame % self.patch_size - pad_num = self.patch_size - r if r else 0 - x = F.pad(x, (0, 0, 0, pad_num), value=0.0) - x = x.reshape(-1, self.patch_size, dim) - x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1) - x = x.reshape(bsz, -1, dim) - return x - - def forward(self, waveform): - x = self.get_frames(waveform) - x = self.fc1(x) - x = self.fc2(x) - x = self.encoder(inputs_embeds=x) - x = x.last_hidden_state - - if self.patch_size != -1: - x = self.pad_patch_insert_cls(x) - x = self.aggregator(inputs_embeds=x) - x = x.last_hidden_state - bsz, _, dim = x.size() - x = x.reshape(-1, self.patch_size + 1, dim) - x = x[:, -1:, :].reshape(bsz, -1, dim) - - x = self.fc3(x) - return x, waveform.unsqueeze(1) - - class AudioVAE(PreTrainedModel): config_class = AudioVAEConfig - def __init__(self, config: AudioVAEConfig): + def __init__(self, config): super().__init__(config) + self.config = config + enc_kwargs = config.enc_kwargs + dec_kwargs = config.dec_kwargs + for key in ("backbone", "input_dim", "latent_dim"): + if key not in enc_kwargs: + raise ValueError(f"AudioVAE.enc_kwargs missing required key: {key}") + for key in ("backbone", "output_dim", "latent_dim"): + if key not in dec_kwargs: + raise ValueError(f"AudioVAE.dec_kwargs missing required key: {key}") + + hop_size = enc_kwargs.get("hop_size", enc_kwargs["input_dim"]) + if enc_kwargs["input_dim"] != hop_size: + raise ValueError(f"AudioVAE encoder input_dim ({enc_kwargs['input_dim']}) != hop_size ({hop_size}).") + if hop_size != dec_kwargs["output_dim"]: + raise ValueError( + f"AudioVAE encoder hop_size ({hop_size}) != decoder output_dim ({dec_kwargs['output_dim']})." + ) self.encoder = Encoder( - encoder_args=config.enc_kwargs["backbone"], - input_dim=config.enc_kwargs["input_dim"], - hop_size=config.enc_kwargs.get("hop_size", 320), - latent_dim=config.enc_kwargs["latent_dim"], + encoder_args=enc_kwargs["backbone"], + input_dim=enc_kwargs["input_dim"], + hop_size=hop_size, + latent_dim=enc_kwargs["latent_dim"], patch_size=config.patch_size, ) self.decoder = Decoder( - decoder_args=config.dec_kwargs["backbone"], - output_dim=config.dec_kwargs["output_dim"], - latent_dim=config.dec_kwargs["latent_dim"], + decoder_args=dec_kwargs["backbone"], + output_dim=dec_kwargs["output_dim"], + latent_dim=dec_kwargs["latent_dim"], patch_size=config.patch_size, ) self.post_init() @@ -367,21 +325,50 @@ class AudioVAE(PreTrainedModel): if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() + @torch.inference_mode() def encode_latent(self, waveform, waveform_length): - from diffusers.models.autoencoders.autoencoder_oobleck import OobleckDiagonalGaussianDistribution + if waveform.ndim != 2: + raise ValueError(f"Expected waveform rank-2 [Batch, Time], got {tuple(waveform.shape)}") + if waveform_length.ndim != 1: + raise ValueError(f"Expected waveform_length rank-1 [Batch], got {tuple(waveform_length.shape)}") + if waveform.shape[0] != waveform_length.shape[0]: + raise ValueError( + "Batch mismatch: " + f"waveform batch={waveform.shape[0]} vs " + f"waveform_length batch={waveform_length.shape[0]}" + ) + if torch.any(waveform_length <= 0): + raise ValueError("waveform_length must be strictly positive.") frame_num = torch.ceil(waveform_length / self.config.enc_kwargs["input_dim"]).to(torch.int32) if self.config.patch_size != -1: frame_num = torch.ceil(frame_num / self.config.patch_size) - h, y = self.encoder(waveform) + h, _ = self.encoder(waveform) h = h.transpose(1, 2) - posterior = OobleckDiagonalGaussianDistribution(h) - latent = posterior.sample() + mean, logvar = torch.chunk(h, 2, dim=1) + logvar = torch.clamp(logvar, -30.0, 20.0) + std = torch.exp(0.5 * logvar) + latent = mean + std * torch.randn_like(mean) latent = latent.transpose(1, 2) return latent, frame_num + @torch.inference_mode() def decode(self, latent, past_key_values=None, use_cache=False, stream_state=(None, None, None), last_chunk=False): + if latent.dim() != 3: + raise ValueError(f"Expected latent rank-3 [B,T,D], got shape={tuple(latent.shape)}") + if latent.shape[0] <= 0: + raise ValueError("latent batch size must be positive.") + + target_dtype = next(self.decoder.parameters()).dtype + target_device = next(self.decoder.parameters()).device + if latent.dtype != target_dtype or latent.device != target_device: + latent = latent.to(device=target_device, dtype=target_dtype) + + expected_latent_dim = self.config.dec_kwargs["latent_dim"] + if latent.shape[-1] != expected_latent_dim: + raise ValueError(f"Latent dim mismatch in decode(): got {latent.shape[-1]}, expected {expected_latent_dim}") + waveform, stream_state, past_key_values = self.decoder.low_level_reconstruct( latent, past_key_values=past_key_values, diff --git a/vllm_omni/model_executor/models/common/ming/dit.py b/vllm_omni/model_executor/models/common/ming/dit.py new file mode 100644 index 000000000..3b4ab1122 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/dit.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +import torch.nn.functional as F +from torch import nn +from x_transformers.x_transformers import apply_rotary_pos_emb + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) + return x + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate="none"): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) + + def forward(self, x): + return self.ff(x) + + +class Attention(nn.Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + qk_norm=None, + pe_attn_head=None, + attn_mask_enabled=True, + ): + super().__init__() + self.dim = dim + self.heads = heads + self.inner_dim = dim_head * heads + self.dropout = dropout + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + if qk_norm is None: + self.q_norm = None + self.k_norm = None + elif qk_norm == "rms_norm": + self.q_norm = RMSNorm(dim_head) + self.k_norm = RMSNorm(dim_head) + else: + raise ValueError(f"Unimplemented qk_norm: {qk_norm}") + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, dim)) + self.to_out.append(nn.Dropout(dropout)) + self.pe_attn_head = pe_attn_head + self.attn_mask_enabled = attn_mask_enabled + + def forward(self, x, mask=None, rope=None): + batch_size = x.shape[0] + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + inner_dim = key.shape[-1] + head_dim = inner_dim // self.heads + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + if self.q_norm is not None: + query = self.q_norm(query) + if self.k_norm is not None: + key = self.k_norm(key) + + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + if self.pe_attn_head is not None: + on = self.pe_attn_head + query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale) + key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + if self.attn_mask_enabled and mask is not None: + valid_sample_indices = mask.any(dim=1) + final_output = torch.zeros_like(query).to(query.device) + attn_mask = mask[valid_sample_indices] + query = query[valid_sample_indices] + key = key[valid_sample_indices] + value = value[valid_sample_indices] + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) + attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2]) + else: + attn_mask = None + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + if self.attn_mask_enabled and mask is not None: + final_output[valid_sample_indices] = x + x = final_output + + x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) + x = x.to(query.dtype) + x = self.to_out[0](x) + x = self.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + + +class DiTBlock(nn.Module): + def __init__( + self, + hidden_size, + num_heads, + mlp_ratio=4.0, + dropout=0.1, + qk_norm=None, + pe_attn_head=None, + attn_mask_enabled=True, + **kwargs, + ): + super().__init__() + self.norm1 = RMSNorm(hidden_size) + self.attn = Attention( + dim=hidden_size, + heads=num_heads, + dim_head=hidden_size // num_heads, + dropout=dropout, + qk_norm=qk_norm, + pe_attn_head=pe_attn_head, + attn_mask_enabled=attn_mask_enabled, + ) + self.norm2 = RMSNorm(hidden_size) + self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh") + + def forward(self, x, mask, rope): + x = x + self.attn(self.norm1(x), mask=mask, rope=rope) + x = x + self.mlp(self.norm2(x)) + return x + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm_final = RMSNorm(hidden_size) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + + def forward(self, x): + x = self.norm_final(x) + x = self.linear(x) + return x + + +class CondEmbedder(nn.Module): + def __init__(self, input_feature_size, hidden_size): + super().__init__() + self.cond_embedder = nn.Linear(input_feature_size, hidden_size) + + def forward(self, llm_cond): + return self.cond_embedder(llm_cond) + + +def get_epss_timesteps(n, device, dtype): + dt = 1 / 32 + predefined_timesteps = { + 5: [0, 2, 4, 8, 16, 32], + 6: [0, 2, 4, 6, 8, 16, 32], + 7: [0, 2, 4, 6, 8, 16, 24, 32], + 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], + 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], + 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], + } + t = predefined_timesteps.get(n, []) + if not t: + return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) + return dt * torch.tensor(t, device=device, dtype=dtype) diff --git a/vllm_omni/model_executor/models/common/ming/fm.py b/vllm_omni/model_executor/models/common/ming/fm.py new file mode 100644 index 000000000..d2cbc6211 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/fm.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from .dit import get_epss_timesteps + + +def apply_sway_sampling(t, sway_sampling_coef): + if sway_sampling_coef is None: + return t + return t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) + + +def build_timesteps(steps, device, dtype, use_epss=True, sway_sampling_coef=-1.0): + if steps <= 0: + raise ValueError(f"steps must be positive, got {steps}") + if use_epss: + t = get_epss_timesteps(steps, device=device, dtype=dtype) + else: + t = torch.linspace(0, 1, steps + 1, device=device, dtype=dtype) + return apply_sway_sampling(t, sway_sampling_coef) + + +class Solver: + def __init__(self, func, y0, sigma=0.25, temperature=1.5) -> None: + self.func = func + self.y0 = y0 + self.sigma = sigma + self.temperature = temperature + + def integrate(self, t): + solution = torch.empty(len(t), *self.y0.shape, dtype=self.y0.dtype, device=self.y0.device) + solution[0] = self.y0 + + j = 1 + y0 = self.y0 + for t0, t1 in zip(t[:-1], t[1:]): + dt = t1 - t0 + f0 = self.func(t0, y0) + dy = dt * f0 + y1 = y0 + dy + + while j < len(t) and t1 >= t[j]: + solution[j] = self._linear_interp(t0, t1, y0, y1, t[j]) + j += 1 + + noise = torch.randn_like(y0) + shift = self.sigma * (self.temperature**0.5) * (abs(dt) ** 0.5) * noise + y0 = y1 + shift + + return solution + + def _linear_interp(self, t0, t1, y0, y1, t): + if t == t0: + return y0 + if t == t1: + return y1 + slope = (t - t0) / (t1 - t0) + return y0 + slope * (y1 - y0) + + +def integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, steps): + for step in range(steps): + dt = t[step + 1] - t[step] + y0 = y0 + fn(t[step], y0) * dt + y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step] + return y0 diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py index 36ff0d456..78c1ea66f 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py @@ -141,24 +141,37 @@ class ConditionalCFM(BASECFM): if isinstance(self.estimator, torch.nn.Module): return self.estimator(x, mask, mu, t, spks, cond) else: + # TensorRT estimator: bind raw device pointers. The flow runs in + # fp32 but the engine may have fp16 I/O (strongly-typed fp16 engine), + # so cast inputs/output to the engine's dtype at the boundary. Keep + # references to the cast buffers alive until execute completes (a bare + # ``.contiguous().data_ptr()`` could free the temp -> dangling ptr). + io_dtype = getattr(self.estimator, "io_dtype", x.dtype) [estimator, stream], trt_engine = self.estimator.acquire_estimator() # NOTE need to synchronize when switching stream torch.cuda.current_stream().synchronize() with stream: - estimator.set_input_shape("x", (2, 80, x.size(2))) - estimator.set_input_shape("mask", (2, 1, x.size(2))) - estimator.set_input_shape("mu", (2, 80, x.size(2))) + x_e = x.to(io_dtype).contiguous() + mask_e = mask.to(io_dtype).contiguous() + mu_e = mu.to(io_dtype).contiguous() + t_e = t.to(io_dtype).contiguous() + spks_e = spks.to(io_dtype).contiguous() + cond_e = cond.to(io_dtype).contiguous() + out_e = torch.empty_like(x_e) + estimator.set_input_shape("x", (2, 80, x_e.size(2))) + estimator.set_input_shape("mask", (2, 1, x_e.size(2))) + estimator.set_input_shape("mu", (2, 80, x_e.size(2))) estimator.set_input_shape("t", (2,)) estimator.set_input_shape("spks", (2, 80)) - estimator.set_input_shape("cond", (2, 80, x.size(2))) + estimator.set_input_shape("cond", (2, 80, x_e.size(2))) data_ptrs = [ - x.contiguous().data_ptr(), - mask.contiguous().data_ptr(), - mu.contiguous().data_ptr(), - t.contiguous().data_ptr(), - spks.contiguous().data_ptr(), - cond.contiguous().data_ptr(), - x.data_ptr(), + x_e.data_ptr(), + mask_e.data_ptr(), + mu_e.data_ptr(), + t_e.data_ptr(), + spks_e.data_ptr(), + cond_e.data_ptr(), + out_e.data_ptr(), ] for i, j in enumerate(data_ptrs): estimator.set_tensor_address(trt_engine.get_tensor_name(i), j) @@ -166,7 +179,7 @@ class ConditionalCFM(BASECFM): assert estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True torch.cuda.current_stream().synchronize() self.estimator.release_estimator(estimator, stream) - return x + return out_e.to(x.dtype) class CausalConditionalCFM(ConditionalCFM): diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index d20312a97..1a387cdcc 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -3,10 +3,16 @@ import os from collections.abc import Iterable, Mapping, Sequence from functools import partial +from math import gcd from threading import Lock +import numpy as np +import onnxruntime import torch import torch.nn as nn +from huggingface_hub import snapshot_download +from scipy.signal import resample_poly +from transformers import Qwen2Config from transformers.feature_extraction_utils import BatchFeature from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions @@ -32,12 +38,15 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct +from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import ( concat_text_with_prompt_ids, extract_speech_feat, - extract_speech_token, extract_spk_embedding, + extract_spk_embedding_trt, extract_text_token, + mel_spectrogram, + unpad_prompt_conditioning, ) from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config @@ -45,6 +54,23 @@ from vllm_omni.utils.speaker_cache import get_speaker_cache logger = init_logger(__name__) +# Process-wide cache of per-model mm-processor runtime components (tokenizer, +# feat_extractor, campplus session/engine). The mm processor is re-created per +# request (mm_processor_cache_gb: 0), so this avoids rebuilding them every time. +_RUNTIME_COMPONENTS_CACHE: dict[str, dict] = {} + + +def _cosyvoice3_trt_enabled() -> bool: + """COSYVOICE3_TRT env toggle (default on) for the optional TensorRT paths. + + Gates both the talker speaker-embedding (campplus) engine and the code2wav + flow-decoder estimator engine. Env-var toggle (cf. mimo_audio's + ``MIMO_AUDIO_TOKENIZER_CUDA_GRAPH``) because these run outside the stage + worker, so deploy-yaml ``hf_overrides`` / per-stage ``env`` do not reach + them — export ``COSYVOICE3_TRT=0`` in the launching shell to disable. + """ + return os.environ.get("COSYVOICE3_TRT", "1") not in ("0", "false", "False", "") + class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo): def get_hf_config(self): @@ -73,39 +99,136 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal if cached_model_dir == model_dir: return - # If model_dir is an HF repo ID (not a local path), resolve to cache - if not os.path.isdir(model_dir): - from huggingface_hub import snapshot_download + # The mm processor is re-created per request (deploy config sets + # ``mm_processor_cache_gb: 0``), so the instance-level guard above never + # hits across requests. Without a process-wide cache, every request + # would re-run ``snapshot_download``, rebuild the Qwen tokenizer and + # create a fresh ONNX campplus session — ~hundreds of ms of pure + # overhead on the TTFP critical path. Build the heavy components once + # per process, keyed by model_dir, and reuse them. + comps = _RUNTIME_COMPONENTS_CACHE.get(model_dir) + if comps is None: + comps = self._build_runtime_components(model_dir, config) + _RUNTIME_COMPONENTS_CACHE[model_dir] = comps + self.tokenizer = comps["tokenizer"] + self.feat_extractor = comps["feat_extractor"] + self.campplus_session = comps["campplus_session"] + self.campplus_trt = comps["campplus_trt"] + self._cached_model_dir = model_dir + self._speaker_cache = get_speaker_cache() + + def _build_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> dict: + """Build the per-model runtime components once (cached process-wide).""" + # If model_dir is an HF repo ID (not a local path), resolve to cache. + if not os.path.isdir(model_dir): model_dir = snapshot_download(model_dir) - import onnxruntime - - from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer - from vllm_omni.model_executor.models.cosyvoice3.utils import mel_spectrogram - - option = onnxruntime.SessionOptions() - option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL - option.intra_op_num_threads = 1 - - self.tokenizer = get_qwen_tokenizer( + tokenizer = get_qwen_tokenizer( token_path=os.path.join(model_dir, config.qwen_pretrain_path), skip_special_tokens=config.skip_special_tokens, version=config.version, ) - self.speech_tokenizer = onnxruntime.InferenceSession( - os.path.join(model_dir, config.speech_tokenizer_path), - sess_options=option, - providers=["CUDAExecutionProvider" if torch.cuda.is_available() else "CPUExecutionProvider"], - ) - self.feat_extractor = partial(mel_spectrogram, **getattr(config, "feat_extractor", {})) - self.campplus_session = onnxruntime.InferenceSession( - os.path.join(model_dir, config.campplus_onxx_path), - sess_options=option, - providers=["CPUExecutionProvider"], - ) - self._cached_model_dir = model_dir - self._speaker_cache = get_speaker_cache() + feat_extractor = partial(mel_spectrogram, **getattr(config, "feat_extractor", {})) + + campplus_onnx_path = os.path.join(model_dir, config.campplus_onxx_path) + # TensorRT speaker-embedding path (default on); a prebuilt TRT engine + # runs campplus on GPU. The engine itself is cached process-wide by + # ``get_campplus_trt``. + campplus_trt = None + if self._speaker_embedding_trt_enabled() and torch.cuda.is_available(): + try: + from vllm_omni.model_executor.models.cosyvoice3.speaker_embedding_trt import ( + get_campplus_trt, + ) + + campplus_trt = get_campplus_trt(campplus_onnx_path, device="cuda") + logger.info("CosyVoice3: using TensorRT campplus speaker embedding") + except Exception as exc: # pragma: no cover - defensive fallback + logger.warning( + "CosyVoice3: TensorRT campplus build failed (%s); falling back to ONNX-Runtime speaker embedding", + exc, + ) + campplus_trt = None + + # Only build the CPU ONNX campplus session when TRT is unavailable — + # otherwise it is never used and creating it (ORT_ENABLE_ALL graph + # optimization) costs ~hundreds of ms. + campplus_session = None + if campplus_trt is None: + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 1 + campplus_session = onnxruntime.InferenceSession( + campplus_onnx_path, + sess_options=option, + providers=["CPUExecutionProvider"], + ) + + return { + "tokenizer": tokenizer, + "feat_extractor": feat_extractor, + "campplus_session": campplus_session, + "campplus_trt": campplus_trt, + } + + @staticmethod + def _speaker_embedding_trt_enabled() -> bool: + """Whether to use the TensorRT campplus speaker-embedding path (default on).""" + return _cosyvoice3_trt_enabled() + + # Class-level cached s3tokenizer model — loaded once per process on first + # call to ``_extract_speech_token_via_s3`` and shared across all + # processor instances. + _s3_model = None + + @classmethod + def _ensure_s3_model(cls): + if cls._s3_model is not None: + return cls._s3_model + + # s3tokenizer is imported lazily (kept off the module top level) so + # callers that don't use the CosyVoice3 talker need not install it. + try: + import s3tokenizer as _s3 + except ImportError as e: + raise ImportError( + "CosyVoice3 speech-token extraction requires the 's3tokenizer' " + "package; install it with `pip install s3tokenizer`." + ) from e + + model = _s3.load_model("speech_tokenizer_v3_25hz") + device = "cuda" if torch.cuda.is_available() else "cpu" + model = model.to(device).eval() + cls._s3_model = (model, _s3, device) + return cls._s3_model + + def _extract_speech_token_via_s3(self, audio, return_device): + """Drop-in replacement for ``extract_speech_token`` that uses the + S3Tokenizer PyTorch model on GPU. Returns the same + ``(speech_token[1, T], speech_token_len[1])`` int32 tensors as the + ONNX path so the rest of ``_call_hf_processor`` is unchanged. + """ + model, _s3, dev = self._ensure_s3_model() + + # audio is a (waveform_ndarray, sr) tuple; resample to 16 kHz mono float32. + wav, sr = audio + wav = np.asarray(wav, dtype=np.float32) + if wav.ndim == 2: + wav = wav.mean(axis=1) + if int(sr) != 16000: + g = gcd(int(sr), 16000) + wav = resample_poly(wav, 16000 // g, int(sr) // g).astype(np.float32) + audio_t = torch.from_numpy(wav) + + mel = _s3.log_mel_spectrogram(audio_t) + mels_p, mels_lens = _s3.padding([mel]) + with torch.inference_mode(): + codes, codes_lens = model.quantize(mels_p.to(dev), mels_lens.to(dev)) + n = int(codes_lens[0].item()) + speech_token = codes[:1, :n].to(dtype=torch.int32, device=return_device) + speech_token_len = torch.tensor([n], dtype=torch.int32, device=return_device) + return speech_token, speech_token_len def _call_hf_processor( self, @@ -185,7 +308,11 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal ) return ft - speech_token, speech_token_len = extract_speech_token(audio, self.speech_tokenizer, device) + # Speech-token extraction via the S3Tokenizer PyTorch model on GPU + # (~30x faster than the bundled ``speech_tokenizer_v3.onnx`` CPU ONNX + # path in this venv). + speech_token, speech_token_len = self._extract_speech_token_via_s3(audio, device) + speech_feat, speech_feat_len = extract_speech_feat(audio, self.feat_extractor, device) if config.sample_rate == 24000: @@ -193,7 +320,10 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal speech_feat, speech_feat_len[:] = speech_feat[:, : 2 * token_len], 2 * token_len speech_token, speech_token_len[:] = speech_token[:, :token_len], token_len - embedding = extract_spk_embedding(audio, self.campplus_session, device) + if self.campplus_trt is not None: + embedding = extract_spk_embedding_trt(audio, self.campplus_trt, device) + else: + embedding = extract_spk_embedding(audio, self.campplus_session, device) # Cache the extracted artifacts for named speakers if cache_key is not None: @@ -318,8 +448,6 @@ class CosyVoice3Model( self.model_stage = vllm_config.model_config.model_stage model_dir = vllm_config.model_config.model if not os.path.isdir(model_dir): - from huggingface_hub import snapshot_download - model_dir = snapshot_download(model_dir) self.model_dir = model_dir self.model = None @@ -377,8 +505,6 @@ class CosyVoice3Model( the pretrained model directory. The cache config is inherited from the parent to enable PagedAttention with the same memory configuration. """ - from transformers import Qwen2Config - qwen_config_path = os.path.join(self.model_dir, self.config.llm["llm"]["pretrain_path"]) qwen_hf_config = Qwen2Config.from_pretrained(qwen_config_path) @@ -469,26 +595,32 @@ class CosyVoice3Model( top_k: int, generator: torch.Generator | None, ) -> int: + """Vectorized nucleus + top-k sampling. + + Distribution-equivalent to the reference iterative implementation: the + keep-set is identical (token i is kept iff + ``cumsum(sorted_probs)[i] - sorted_probs[i] < top_p`` AND ``i < top_k``) + and the renormalized sampling distribution matches, but the exact token + drawn for a given seed is NOT guaranteed to match. The reference draws + via ``multinomial`` over the stacked kept subset while this draws over + the full sorted vector (zeroed outside the keep-set), so the generator + advances over different-sized inputs and may yield a different sample. + The win: no per-token ``.item()`` D2H syncs from the Python loop — + those dominated the sampler CPU time in profiling. + """ probs = weighted_scores.softmax(dim=0) sorted_prob, sorted_idx = probs.sort(descending=True, stable=True) - kept_probs: list[torch.Tensor] = [] - kept_indices: list[torch.Tensor] = [] - cum_prob = 0.0 - max_keep = len(sorted_idx) if top_k <= 0 else min(int(top_k), len(sorted_idx)) - for i in range(len(sorted_idx)): - if cum_prob < top_p and len(kept_probs) < max_keep: - cum_prob += float(sorted_prob[i].item()) - kept_probs.append(sorted_prob[i]) - kept_indices.append(sorted_idx[i]) - else: - break - - if not kept_probs: - return int(sorted_idx[0].item()) - - sample_probs = torch.stack(kept_probs) - sample_idx = cls._multinomial_sample(sample_probs, generator=generator) - return int(torch.stack(kept_indices)[int(sample_idx.item())].item()) + cum_before = sorted_prob.cumsum(dim=0) - sorted_prob + mask = cum_before < top_p + if top_k > 0: + n = sorted_prob.shape[0] + mask = mask & (torch.arange(n, device=mask.device) < min(int(top_k), n)) + weights = sorted_prob * mask.to(sorted_prob.dtype) + # First token always passes (cum_before[0] = 0 < top_p for any top_p > 0), + # so ``weights`` is guaranteed to have at least one nonzero entry. The + # final ``.item()`` is the ONLY D2H sync per call. + sample_idx = torch.multinomial(weights, 1, replacement=True, generator=generator) + return int(sorted_idx[sample_idx].item()) @classmethod def _ras_sample_one( @@ -627,11 +759,21 @@ class CosyVoice3Model( else: raise RuntimeError(f"compute_logits is only valid for {self.model_stage}.") - def embed_multimodal(self, **kwargs: object) -> torch.Tensor: + def embed_multimodal(self, **kwargs: object): if self.model_stage == "cosyvoice3_talker": speech_token = kwargs["speech_token"] - speech_token_emb = self.model.speech_embedding(speech_token) - return speech_token_emb + # vLLM's _execute_mm_encoder batches ALL multimodal items scheduled + # in one engine step into a single call, expecting one embedding + # tensor per item back. When >=2 requests prefill in the same step + # (likely once mm preprocessing is fast, e.g. TensorRT speaker + # embedding), speech_token arrives as a list of per-item tensors + # (variable length), so embed it per item and return a list of + # [T_i, emb] tensors. The single-item path keeps returning a + # [1, T, emb] tensor (the caller's extend() iterates dim 0). + if isinstance(speech_token, (list, tuple)): + emb_dim = self.model.speech_embedding.weight.shape[1] + return [self.model.speech_embedding(t).reshape(-1, emb_dim) for t in speech_token] + return self.model.speech_embedding(speech_token) else: raise RuntimeError(f"embed_multimodal is only valid for {self.model_stage}.") @@ -720,6 +862,121 @@ class CosyVoice3Model( else: raise RuntimeError(f"embed_input_ids is not valid for {self.model_stage}.") + @staticmethod + def _split_prompt_conditioning(speech_token, speech_feat, embedding, speech_token_len): + """Split collated prompt conditioning into per-request lists. + + Returns ``(speech_token_list, speech_feat_list, embedding_list, + speech_token_len_list)``, one rank-correct tensor per request: + ``speech_token`` ``[1, T]`` (possibly right-padded), ``speech_feat`` + ``[1, 2T, F]``, ``embedding`` ``[1, D]``, ``speech_token_len`` ``[1, 1]``. + + This runs inside the talker ``forward`` which may be CUDA-graph + captured, so it must NOT trigger any host<->device sync (no ``.item()`` + / ``.tolist()`` / Python-int slicing by tensor value). Splitting is pure + on-device indexing; right-padding is dropped later in the (eager) + code2wav stage using the per-request ``speech_token_len``. Robust to + inputs arriving as padded ``[B, ...]`` batch tensors or per-request + lists. + """ + + def _rows(x, want_dim): + if isinstance(x, (list, tuple)): + items = list(x) + elif isinstance(x, torch.Tensor): + # ``x.shape[0]`` is a static Python int — no device sync. + items = [x[i] for i in range(x.shape[0])] + else: + return None + out = [] + for t in items: + if isinstance(t, torch.Tensor): + while t.dim() < want_dim: + t = t.unsqueeze(0) + if t.shape[0] != 1: + t = t[:1] + t = t.contiguous() + out.append(t) + return out + + st_out = _rows(speech_token, 2) + sf_out = _rows(speech_feat, 3) + emb_out = _rows(embedding, 2) + stl_out = _rows(speech_token_len, 2) + return st_out, sf_out, emb_out, stl_out + + def _resolve_flow_estimator_onnx(self) -> str | None: + """Locate the flow-decoder estimator ONNX for the TensorRT engine. + + Prefers the fp16 (strongly-typed) ONNX → fp16 engine. Order: env + ``COSYVOICE3_ESTIMATOR_ONNX`` → ``/`` → fetched + from ``flow_estimator_onnx_repo`` → bundled fp32 ONNX (fp32+TF32). Each + is checked for existence; returns ``None`` if nothing is found. + """ + env_path = os.environ.get("COSYVOICE3_ESTIMATOR_ONNX") + if env_path and os.path.exists(env_path): + return env_path + + fp16_name = getattr(self.config, "flow_estimator_onnx_path", "flow.decoder.estimator.autocast_fp16.onnx") + local_fp16 = os.path.join(self.model_dir, fp16_name) + if os.path.exists(local_fp16): + return local_fp16 + + repo = getattr(self.config, "flow_estimator_onnx_repo", None) + if repo: + try: + fetched_dir = snapshot_download(repo, allow_patterns=[fp16_name]) + fetched = os.path.join(fetched_dir, fp16_name) + if os.path.exists(fetched): + return fetched + except Exception as exc: # pragma: no cover - network/repo issues + logger.warning("CosyVoice3 code2wav: could not fetch fp16 estimator ONNX from %s (%s)", repo, exc) + + fp32_name = getattr(self.config, "flow_estimator_onnx_path_fp32", "flow.decoder.estimator.fp32.onnx") + local_fp32 = os.path.join(self.model_dir, fp32_name) + if os.path.exists(local_fp32): + logger.info("CosyVoice3 code2wav: fp16 estimator ONNX unavailable, falling back to fp32 (%s)", local_fp32) + return local_fp32 + return None + + def _maybe_enable_code2wav_trt(self) -> None: + """Swap the flow-decoder estimator to a TensorRT engine once (lazy). + + Runs on the first code2wav step — after weights are loaded — so the + torch estimator is fully built first and then dropped. No-op unless + ``COSYVOICE3_TRT`` is on (default), CUDA is available, and the estimator + ONNX ships with the model. Falls back to the torch estimator on any + failure. The upstream ``CausalConditionalCFM.forward_estimator`` detects + the non-``nn.Module`` estimator and drives the TRT engine. + """ + if getattr(self, "_code2wav_trt_done", False): + return + self._code2wav_trt_done = True + if not (_cosyvoice3_trt_enabled() and torch.cuda.is_available()): + return + onnx_path = self._resolve_flow_estimator_onnx() + if onnx_path is None: + logger.warning("CosyVoice3 code2wav: no flow-estimator ONNX available; keeping torch estimator") + return + try: + from vllm_omni.model_executor.models.cosyvoice3.flow_estimator_trt import ( + build_flow_estimator_trt, + ) + + wrapper = build_flow_estimator_trt(onnx_path, device="cuda") + # ``estimator`` is a registered nn.Module submodule; delete it first + # (frees the torch estimator weights) so the TRT wrapper can be set + # as a plain attribute — nn.Module.__setattr__ rejects non-Modules. + decoder = self.code2wav.flow_model.decoder + del decoder.estimator + decoder.estimator = wrapper + logger.info("CosyVoice3: using TensorRT flow-decoder estimator (code2wav)") + except Exception as exc: # pragma: no cover - defensive fallback + logger.warning( + "CosyVoice3 code2wav: TensorRT estimator build failed (%s); keeping torch estimator", + exc, + ) + def forward( self, input_ids: torch.Tensor, @@ -741,18 +998,43 @@ class CosyVoice3Model( if "speech_token" in kwargs: # Prompt conditioning tensors for code2wav: live under # ``embed.*`` per OmniPayloadStruct schema. + # + # vLLM hands these mm fields to forward() as collated, padded + # batch tensors (speech_token [B, maxT], speech_feat + # [B, 2*maxT, F], embedding [B, D]). Emitting them raw makes the + # downstream per-request payload split fragile at batch>1: it + # intermittently de-batches speech_token to 1-D and leaks the + # whole [B, D] embedding to every request, which corrupts voice + # conditioning and crashes code2wav (`prompt_token.shape[1]`). + # Instead, split into an explicit per-request list of unpadded, + # correctly-ranked tensors so ``to_payload_element`` splits them + # deterministically by request index (list[idx]). + speech_token_list, speech_feat_list, embedding_list, speech_token_len_list = ( + self._split_prompt_conditioning( + kwargs.get("speech_token"), + kwargs.get("speech_feat"), + kwargs.get("embedding"), + kwargs.get("speech_token_len"), + ) + ) multimodal_outputs = to_dict( OmniPayloadStruct( embed=EmbeddingsStruct( - speech_token=kwargs.get("speech_token"), - speech_feat=kwargs.get("speech_feat"), - embedding=kwargs.get("embedding"), + speech_token=speech_token_list, + speech_feat=speech_feat_list, + speech_token_len=speech_token_len_list, + embedding=embedding_list, ), ) ) return OmniOutput(text_hidden_states=hidden_states, multimodal_outputs=multimodal_outputs) elif self.model_stage == "cosyvoice3_code2wav": + # Lazily swap the flow-decoder estimator to a TensorRT engine on the + # first code2wav step (after weights are loaded), gated by the same + # COSYVOICE3_TRT env toggle as the talker speaker embedding. + self._maybe_enable_code2wav_trt() + runtime_info = kwargs.get("model_intermediate_buffer") if runtime_info is None: runtime_info = kwargs.get("runtime_additional_information", []) @@ -784,6 +1066,11 @@ class CosyVoice3Model( speech_token = embed.speech_token if embed else None speech_feat = embed.speech_feat if embed else None embedding = embed.embedding if embed else None + # Drop any right-padding carried from batched talker emission. + if speech_token is not None and speech_feat is not None: + speech_token, speech_feat = unpad_prompt_conditioning( + speech_token, speech_feat, embed.speech_token_len if embed else None + ) if speech_token is None or speech_feat is None or embedding is None: if stream_finished and req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"): with self._stream_audio_cache_lock: diff --git a/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py b/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py new file mode 100644 index 000000000..230fb14fb --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TensorRT engine for the CosyVoice3 flow-decoder (CFM) DiT estimator. + +The estimator is the per-step network the conditional flow-matching ODE solver +calls during code2wav (token -> mel). It dominates code2wav latency; the +upstream ``CausalConditionalCFM.forward_estimator`` already supports running it +through a TensorRT engine (it switches on ``self.estimator`` not being an +``nn.Module`` and drives it via ``acquire_estimator`` / ``execute_async_v3``). +This module builds that engine from the bundled ``flow.decoder.estimator*.onnx`` +and wraps it so it can be dropped in for the torch estimator. + +The estimator engine has 6 inputs ``x, mask, mu, t, spks, cond`` and one output. +``x/mask/mu/cond`` carry a dynamic time dim; ``t``/``spks`` are fixed, so only +the former need an optimization profile (TRT infers the latter from the ONNX). +Shapes mirror the CosyVoice runtime. + +Precision: TensorRT >= 11 dropped the weakly-typed FP16/INT8 builder flags, so +fp16 only comes from a STRONGLY_TYPED network built from an fp16 ONNX +(``*autocast_fp16*``, fp16 I/O). An fp32 ONNX is built fp32 + the TF32 matmul +flag. ``EXPLICIT_BATCH`` is implicit (no flag) and ``ITensor.dtype`` is +read-only, so neither is set here. +""" + +from __future__ import annotations + +import os +import queue + +import torch +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.cosyvoice3.speaker_embedding_trt import ( + _resolve_plan_path, + _trt_logger, +) + +logger = init_logger(__name__) + +# Optimization-profile shapes for the dynamic-length inputs (CFG batch dim = 2, +# 80 mel channels, time dim min/opt/max). Matches CosyVoice's token2wav. +_DYNAMIC_INPUTS = ("x", "mask", "mu", "cond") +_MIN_SHAPES = ((2, 80, 4), (2, 1, 4), (2, 80, 4), (2, 80, 4)) +_OPT_SHAPES = ((2, 80, 500), (2, 1, 500), (2, 80, 500), (2, 80, 500)) +_MAX_SHAPES = ((2, 80, 3000), (2, 1, 3000), (2, 80, 3000), (2, 80, 3000)) + + +def _is_fp16_onnx(onnx_path: str) -> bool: + """Heuristic: the project's fp16 estimator ONNX is exported strongly-typed + and named ``*autocast_fp16*`` / ``*fp16*`` (vs ``*fp32*``).""" + name = os.path.basename(onnx_path).lower() + return "fp16" in name or "autocast" in name + + +def _convert_onnx_to_trt(onnx_path: str, plan_path: str, strongly_typed: bool) -> None: + import tensorrt as trt + + logger.info( + "Building flow-estimator TensorRT engine from %s (%s) ...", + onnx_path, + "strongly-typed/fp16" if strongly_typed else "fp32+TF32", + ) + trt_logger = _trt_logger() + builder = trt.Builder(trt_logger) + # STRONGLY_TYPED takes precision from the ONNX graph (fp16 engine from an + # fp16 ONNX) — this is the only way to get fp16 in TRT>=11, which dropped + # the weakly-typed FP16 BuilderFlag. Otherwise EXPLICIT_BATCH is implicit, + # so create the network with no flags. + if strongly_typed: + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + else: + network = builder.create_network(0) + parser = trt.OnnxParser(network, trt_logger) + with open(onnx_path, "rb") as f: + if not parser.parse(f.read()): + errs = "; ".join(str(parser.get_error(i)) for i in range(parser.num_errors)) + raise ValueError(f"Failed to parse {onnx_path}: {errs}") + + config = builder.create_builder_config() + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 32) # 4 GiB + if not strongly_typed: + # fp32 ONNX: enable the best available reduced-precision matmul flag + # (FP16 on older TRT, else TF32 on Ampere+/Hopper). For a strongly-typed + # network the precision is fixed by the graph, so no flag is set. + for _flag_name in ("FP16", "TF32"): + _flag = getattr(trt.BuilderFlag, _flag_name, None) + if _flag is not None: + config.set_flag(_flag) + break + + profile = builder.create_optimization_profile() + for name, mn, op, mx in zip(_DYNAMIC_INPUTS, _MIN_SHAPES, _OPT_SHAPES, _MAX_SHAPES): + profile.set_shape(name, mn, op, mx) + config.add_optimization_profile(profile) + + engine_bytes = builder.build_serialized_network(network, config) + if engine_bytes is None: + raise RuntimeError(f"TensorRT failed to build flow-estimator engine from {onnx_path}") + tmp = plan_path + ".tmp" + with open(tmp, "wb") as f: + f.write(engine_bytes) + os.replace(tmp, plan_path) + logger.info("Wrote flow-estimator TensorRT engine to %s", plan_path) + + +class TrtContextWrapper: + """Pool of TensorRT execution contexts for the flow estimator. + + Exposes the ``acquire_estimator`` / ``release_estimator`` contract that + ``CausalConditionalCFM.forward_estimator`` expects. + """ + + def __init__( + self, engine, device: str | torch.device, io_dtype: torch.dtype = torch.float32, trt_concurrent: int = 1 + ): + self.trt_engine = engine + # Engine I/O dtype (fp16 for a strongly-typed fp16 engine). The flow runs + # in fp32, so forward_estimator casts to/from this at the boundary. + self.io_dtype = io_dtype + self._pool: queue.Queue = queue.Queue(maxsize=trt_concurrent) + for _ in range(trt_concurrent): + ctx = engine.create_execution_context() + assert ctx is not None, "failed to create TRT execution context (out of memory?)" + stream = torch.cuda.stream(torch.cuda.Stream(torch.device(device))) + self._pool.put([ctx, stream]) + + def acquire_estimator(self): + return self._pool.get(), self.trt_engine + + def release_estimator(self, context, stream): + self._pool.put([context, stream]) + + +def build_flow_estimator_trt(onnx_path: str, device: str | torch.device) -> TrtContextWrapper: + """Build/load the flow-estimator TRT engine and return a context-pool wrapper. + + An fp16 ONNX (``*autocast_fp16*``) is built as a strongly-typed network (the + only way to get fp16 in TRT>=11); an fp32 ONNX is built fp32 + TF32. + """ + import tensorrt as trt + + strongly_typed = _is_fp16_onnx(onnx_path) + plan_path = _resolve_plan_path(onnx_path, prefix="flow_estimator") + if not os.path.exists(plan_path) or os.path.getsize(plan_path) == 0: + _convert_onnx_to_trt(onnx_path, plan_path, strongly_typed=strongly_typed) + + runtime = trt.Runtime(_trt_logger()) + with open(plan_path, "rb") as f: + engine = runtime.deserialize_cuda_engine(f.read()) + if engine is None: + raise RuntimeError(f"Failed to deserialize flow-estimator TensorRT engine {plan_path}") + logger.info("Loaded flow-estimator TensorRT engine (%s)", plan_path) + io_dtype = torch.float16 if strongly_typed else torch.float32 + return TrtContextWrapper(engine, device=device, io_dtype=io_dtype) diff --git a/vllm_omni/model_executor/models/cosyvoice3/speaker_embedding_trt.py b/vllm_omni/model_executor/models/cosyvoice3/speaker_embedding_trt.py new file mode 100644 index 000000000..61aacd917 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/speaker_embedding_trt.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TensorRT speaker-embedding (campplus) extractor for CosyVoice3. + +Drop-in accelerator for the ONNX-Runtime campplus session used in the +multimodal processor. ``campplus.onnx`` maps an fbank feature +``[B, T, 80]`` to a 192-d speaker embedding ``[B, 192]``; on this stack the +ONNX path only has a CPU execution provider, so it dominates per-request +prompt-processing latency. Building a TensorRT engine once and running it on +GPU removes that bottleneck. + +The engine is built lazily from ``campplus.onnx`` on first use and cached as +a ``.plan`` file (engines are device/TRT-version specific, so the cache key +includes both). Set ``COSYVOICE3_TRT_CACHE`` to override the cache dir. + +Written against TensorRT >= 10/11: ``EXPLICIT_BATCH`` is the implicit default +(create the network with no flags) and ``ITensor.dtype`` is read-only (I/O +dtypes come from the ONNX graph), so neither is set here. +""" + +from __future__ import annotations + +import hashlib +import os +import threading + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# campplus fbank sequence length bounds for the optimization profile. +_MIN_T = 4 +_OPT_T = 500 +_MAX_T = 3000 +_EMBED_DIM = 192 + + +def _trt_logger(): + import tensorrt as trt + + return trt.Logger(trt.Logger.WARNING) + + +def _convert_onnx_to_trt(onnx_path: str, plan_path: str) -> None: + import tensorrt as trt + + logger.info("Building campplus TensorRT engine from %s ...", onnx_path) + trt_logger = _trt_logger() + builder = trt.Builder(trt_logger) + # EXPLICIT_BATCH is implicit in TRT>=10; create the network with no flags. + network = builder.create_network(0) + parser = trt.OnnxParser(network, trt_logger) + with open(onnx_path, "rb") as f: + if not parser.parse(f.read()): + errs = "; ".join(str(parser.get_error(i)) for i in range(parser.num_errors)) + raise ValueError(f"Failed to parse {onnx_path}: {errs}") + + config = builder.create_builder_config() + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 31) # 2 GiB + + input_name = network.get_input(0).name + profile = builder.create_optimization_profile() + profile.set_shape( + input_name, + (1, _MIN_T, 80), + (1, _OPT_T, 80), + (1, _MAX_T, 80), + ) + config.add_optimization_profile(profile) + + engine_bytes = builder.build_serialized_network(network, config) + if engine_bytes is None: + raise RuntimeError(f"TensorRT failed to build campplus engine from {onnx_path}") + tmp = plan_path + ".tmp" + with open(tmp, "wb") as f: + f.write(engine_bytes) + os.replace(tmp, plan_path) + logger.info("Wrote campplus TensorRT engine to %s", plan_path) + + +def _resolve_plan_path(onnx_path: str, prefix: str = "campplus") -> str: + cache_dir = os.environ.get("COSYVOICE3_TRT_CACHE") or os.path.join( + os.path.expanduser("~"), ".cache", "vllm_omni", "cosyvoice3_trt" + ) + os.makedirs(cache_dir, exist_ok=True) + try: + # No-arg form targets the current device (avoids the banned + # torch.cuda.current_device per ruff TID251). + dev_name = torch.cuda.get_device_name() + except Exception: + dev_name = "unknown" + import tensorrt as trt + + st = os.stat(onnx_path) + key = f"{os.path.abspath(onnx_path)}|{st.st_size}|{int(st.st_mtime)}|{dev_name}|trt{trt.__version__}" + digest = hashlib.sha1(key.encode()).hexdigest()[:16] + return os.path.join(cache_dir, f"{prefix}_{digest}.plan") + + +class CampplusTRT: + """Lazily-built TensorRT campplus speaker-embedding extractor.""" + + def __init__(self, onnx_path: str, device: str | torch.device): + import tensorrt as trt + + self.device = torch.device(device) + plan_path = _resolve_plan_path(onnx_path) + if not os.path.exists(plan_path) or os.path.getsize(plan_path) == 0: + _convert_onnx_to_trt(onnx_path, plan_path) + + runtime = trt.Runtime(_trt_logger()) + with open(plan_path, "rb") as f: + self.engine = runtime.deserialize_cuda_engine(f.read()) + if self.engine is None: + raise RuntimeError(f"Failed to deserialize campplus TensorRT engine {plan_path}") + self.context = self.engine.create_execution_context() + self.input_name = self.engine.get_tensor_name(0) + self.output_name = self.engine.get_tensor_name(1) + # TRT execution contexts are not safe for concurrent enqueue; serialize. + self._lock = threading.Lock() + logger.info("Loaded campplus TensorRT engine (%s)", plan_path) + + @torch.inference_mode() + def __call__(self, feat: torch.Tensor) -> torch.Tensor: + """feat: ``[T, 80]`` fbank features. Returns ``[1, 192]`` embedding.""" + x = feat.unsqueeze(0).to(self.device, dtype=torch.float32).contiguous() + out = torch.empty((x.shape[0], _EMBED_DIM), device=self.device, dtype=torch.float32) + stream = torch.cuda.current_stream(self.device) + with self._lock: + self.context.set_input_shape(self.input_name, tuple(x.shape)) + self.context.set_tensor_address(self.input_name, x.data_ptr()) + self.context.set_tensor_address(self.output_name, out.data_ptr()) + ok = self.context.execute_async_v3(stream.cuda_stream) + if not ok: + raise RuntimeError("campplus TensorRT execute_async_v3 failed") + stream.synchronize() + return out + + +# Process-wide cache: the mm processor is re-created per request (the deploy +# config disables the mm-processor object cache), so building a fresh +# CampplusTRT each time would re-deserialize the ~33 MB engine and allocate a +# new execution context on every request — inflating TTFP. Cache by +# (engine path, device) and reuse; the context serializes concurrent enqueues +# via its own lock, so a single shared instance is safe across requests. +_CAMPPLUS_CACHE: dict[tuple[str, str], CampplusTRT] = {} + + +def get_campplus_trt(onnx_path: str, device: str | torch.device) -> CampplusTRT: + """Return a process-wide cached :class:`CampplusTRT` (engine + context).""" + key = (os.path.abspath(onnx_path), str(device)) + inst = _CAMPPLUS_CACHE.get(key) + if inst is None: + inst = CampplusTRT(onnx_path, device) + _CAMPPLUS_CACHE[key] = inst + return inst diff --git a/vllm_omni/model_executor/models/cosyvoice3/utils.py b/vllm_omni/model_executor/models/cosyvoice3/utils.py index 0bf0cccb1..5827e6bbb 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/utils.py +++ b/vllm_omni/model_executor/models/cosyvoice3/utils.py @@ -215,6 +215,44 @@ def extract_spk_embedding(prompt_wav, campplus_session, device): return embedding +def unpad_prompt_conditioning(speech_token, speech_feat, speech_token_len): + """Drop right-padding from per-request prompt conditioning. + + The talker emits ``speech_token`` ``[1, T]`` possibly right-padded to the + batch max, plus the true ``speech_token_len``; ``speech_feat`` is at the 2:1 + mel:token ratio (``[1, 2T, F]``). This trims both to the real length. Called + in the (eager) code2wav stage where a host read of the length is allowed. + Returns ``(speech_token, speech_feat)`` unchanged when length is unavailable. + """ + if not isinstance(speech_token, torch.Tensor) or speech_token_len is None: + return speech_token, speech_feat + try: + token_len = int(torch.as_tensor(speech_token_len).reshape(-1)[0].item()) + except (RuntimeError, ValueError, IndexError): + return speech_token, speech_feat + if token_len <= 0 or speech_token.dim() < 2: + return speech_token, speech_feat + token_len = min(token_len, int(speech_token.shape[1])) + speech_token = speech_token[:, :token_len] + if isinstance(speech_feat, torch.Tensor) and speech_feat.dim() >= 2: + speech_feat = speech_feat[:, : 2 * token_len] + return speech_token, speech_feat + + +def extract_spk_embedding_trt(prompt_wav, campplus_trt, device): + """TensorRT counterpart of ``extract_spk_embedding``. + + Identical fbank front-end; the campplus forward runs on a prebuilt + TensorRT engine (GPU) instead of the CPU ONNX-Runtime session. Returns the + same ``[1, 192]`` embedding tensor on ``device``. + """ + speech = load_wav(prompt_wav, 16000) + feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000) + feat = feat - feat.mean(dim=0, keepdim=True) + embedding = campplus_trt(feat) + return embedding.to(device) + + def extract_text_token(text, tokenizer, allowed_special): text_token = tokenizer.encode(text, allowed_special=allowed_special) text_token = torch.tensor([text_token], dtype=torch.int32) diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py index a51d692b9..cacf82ee5 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py @@ -162,6 +162,54 @@ def _build_delay_pattern(codes: torch.Tensor) -> torch.Tensor: _ENCODER_CACHE: Any | None = None _K2_OMNIVOICE_REPO = "k2-fsa/OmniVoice" _K2_OMNIVOICE_SUBDIR = "audio_tokenizer" +_AUDIO_TOKENIZER_PATH_ENVS = ( + "HIGGS_AUDIO_TOKENIZER_PATH", + "HIGGS_AUDIO_V2_TOKENIZER_PATH", +) + + +def _normalize_audio_tokenizer_dir(path: str) -> str | None: + if not path: + return None + expanded = os.path.abspath(os.path.expanduser(path)) + if os.path.isfile(os.path.join(expanded, "config.json")): + return expanded + nested = os.path.join(expanded, _K2_OMNIVOICE_SUBDIR) + if os.path.isfile(os.path.join(nested, "config.json")): + return nested + return None + + +def _resolve_audio_tokenizer_dir() -> str | None: + """Resolve the Higgs codec encoder dir without requiring online HF access. + + Voice cloning runs in the API server process. Production H20 jobs are often + offline, so prefer explicit local directories and already-populated HF cache + before falling back to ``snapshot_download``. + """ + for env_name in _AUDIO_TOKENIZER_PATH_ENVS: + candidate = _normalize_audio_tokenizer_dir(os.getenv(env_name, "")) + if candidate is not None: + return candidate + + from huggingface_hub import try_to_load_from_cache + + cached_config = try_to_load_from_cache( + repo_id=_K2_OMNIVOICE_REPO, + filename=f"{_K2_OMNIVOICE_SUBDIR}/config.json", + ) + if isinstance(cached_config, str) and os.path.isfile(cached_config): + return os.path.dirname(cached_config) + + from huggingface_hub.constants import HF_HUB_CACHE + + safe = _K2_OMNIVOICE_REPO.replace("/", "--") + snapshots_dir = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots") + if os.path.isdir(snapshots_dir): + for rev in os.listdir(snapshots_dir): + candidate = _normalize_audio_tokenizer_dir(os.path.join(snapshots_dir, rev)) + if candidate is not None: + return candidate def _load_audio_tokenizer(): @@ -180,14 +228,17 @@ def _load_audio_tokenizer(): Only the ``audio_tokenizer/`` subdirectory (~806 MB) is downloaded. """ - from huggingface_hub import snapshot_download from transformers import HiggsAudioV2TokenizerModel - repo_path = snapshot_download( - _K2_OMNIVOICE_REPO, - allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], - ) - audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR) + audio_tokenizer_dir = _resolve_audio_tokenizer_dir() + if audio_tokenizer_dir is None: + from huggingface_hub import snapshot_download + + repo_path = snapshot_download( + _K2_OMNIVOICE_REPO, + allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], + ) + audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR) device = "cuda" if torch.cuda.is_available() else "cpu" model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device) return model.eval() diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py b/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py new file mode 100644 index 000000000..052d51eef --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vllm-omni integration for boson-ai's higgs-audio v3 (two-stage TTS).""" + +from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, +) + +__all__ = ["HiggsAudioV3Config"] diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py b/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py new file mode 100644 index 000000000..f6db731c7 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration class for higgs-audio v3 (HiggsMultimodalQwen3) in vllm-omni. + +``HiggsAudioV3Config.from_pretrained(model_path)`` returns a config with +``tts_token_id``, ``text_token_id``, ``audio_continuation_id``, and +``eos_token_id`` already resolved from the checkpoint tokenizer. If the +tokenizer is unavailable or missing required specials, the load raises. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any + +import transformers +from transformers import AutoConfig, PretrainedConfig + +__all__ = ["HiggsAudioV3Config"] + +logger = logging.getLogger(__name__) + +_QWEN3_ROPE_THETA = 1_000_000 + +_REQUIRED_SPECIALS = ("<|tts|>", "<|text|>", "<|audio|>") + + +def _build_text_config(raw: Any) -> PretrainedConfig: + if isinstance(raw, PretrainedConfig): + return raw + cfg = dict(raw or {}) + model_type = cfg.get("model_type", "qwen3") + if model_type == "qwen3" and cfg.get("rope_theta") is None: + cfg["rope_theta"] = _QWEN3_ROPE_THETA + try: + cfg_cls = transformers.CONFIG_MAPPING[model_type] + except KeyError as exc: + raise ValueError(f"Unknown text backbone model_type {model_type!r}") from exc + return cfg_cls(**cfg) + + +def _resolve_model_dir(pretrained_model_name_or_path: str) -> str | None: + """Resolve a model name/path to a local directory containing tokenizer files.""" + if os.path.isdir(pretrained_model_name_or_path): + return pretrained_model_name_or_path + try: + from huggingface_hub import try_to_load_from_cache + + cached = try_to_load_from_cache(repo_id=pretrained_model_name_or_path, filename="tokenizer_config.json") + if isinstance(cached, str) and os.path.isfile(cached): + return os.path.dirname(cached) + except Exception: + pass + + try: + from huggingface_hub.constants import HF_HUB_CACHE + + safe = pretrained_model_name_or_path.replace("/", "--") + snapshots = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots") + if os.path.isdir(snapshots): + for rev in os.listdir(snapshots): + candidate = os.path.join(snapshots, rev) + if os.path.isfile(os.path.join(candidate, "tokenizer_config.json")): + return candidate + except Exception: + pass + return None + + +class HiggsAudioV3Config(PretrainedConfig): + """Typed config for higgs-audio v3 (HiggsMultimodalQwen3). + + ``from_pretrained()`` automatically resolves ``<|tts|>``, ``<|text|>``, + ``<|audio|>`` and ``eos_token_id`` from the checkpoint tokenizer. + """ + + model_type: str = "higgs_multimodal_qwen3" + is_composition = True + + def __init__( + self, + audio_encoder_config: dict[str, Any] | None = None, + text_config: dict[str, Any] | PretrainedConfig | None = None, + audio_token_id: int = -100, + mel_per_sample: int = 8, + num_codebooks: int = 8, + codebook_size: int = 1026, + audio_stream_bos_id: int = 1024, + audio_stream_eos_id: int = 1025, + sample_rate: int = 24000, + frame_rate: int = 25, + tts_token_id: int | None = None, + text_token_id: int | None = None, + audio_continuation_id: int | None = None, + enable_flashinfer_api_unwrap: bool = True, + enable_mlp_cudagraph: bool = True, + **kwargs: Any, + ) -> None: + # Legacy perf knob removed: Higgs v3 scheduler tokens now come from + # SamplerOutput.sampled_token_ids only. + kwargs.pop("enable_cpu_token_override", None) + self.audio_token_id = audio_token_id + self.mel_per_sample = mel_per_sample + + if audio_encoder_config is None: + audio_encoder_config = { + "encoder_type": "discrete", + "num_codebooks": num_codebooks, + "vocab_size": codebook_size, + "tie_word_embeddings": True, + } + self.audio_encoder_config = audio_encoder_config + + self.num_codebooks = int(audio_encoder_config.get("num_codebooks", num_codebooks)) + if self.num_codebooks <= 0: + raise ValueError(f"num_codebooks must be > 0, got {self.num_codebooks}") + self.codebook_size = int(audio_encoder_config.get("vocab_size", codebook_size)) + self.tie_modality_embeddings = bool(audio_encoder_config.get("tie_word_embeddings", True)) + + self.audio_stream_bos_id = audio_stream_bos_id + self.audio_stream_eos_id = audio_stream_eos_id + self.sample_rate = sample_rate + self.frame_rate = frame_rate + + self.text_config = _build_text_config(text_config) + self.audio_hidden_size = int(audio_encoder_config.get("out_dim", self.text_config.hidden_size)) + + self.tts_token_id = tts_token_id + self.text_token_id = text_token_id + self.audio_continuation_id = audio_continuation_id + self.enable_flashinfer_api_unwrap = bool(enable_flashinfer_api_unwrap) + self.enable_mlp_cudagraph = bool(enable_mlp_cudagraph) + + super().__init__(**kwargs) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs: Any) -> HiggsAudioV3Config: + """Load config and resolve special token IDs from the checkpoint tokenizer. + + Passes the original ``pretrained_model_name_or_path`` (local dir or + HF repo ID) directly to ``AutoTokenizer.from_pretrained()`` so it can + handle cache hits, downloads, and local paths uniformly. Raises if the + tokenizer is missing required specials. + """ + config = super().from_pretrained(pretrained_model_name_or_path, **kwargs) + config.resolve_special_tokens(pretrained_model_name_or_path) + return config + + def get_text_config(self, decoder: bool = False) -> PretrainedConfig: + del decoder + return self.text_config + + @property + def num_real_codes(self) -> int: + return self.audio_stream_bos_id + + @property + def hidden_size(self) -> int: + return self.text_config.hidden_size + + def resolve_special_tokens(self, model_path: str) -> None: + """Resolve <|tts|>, <|text|>, <|audio|> and eos from the HF tokenizer. + + Raises ``ValueError`` if any of the 3 required specials is missing + from the tokenizer's added vocabulary. + """ + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + vocab = dict(tokenizer.get_added_vocab()) + + missing = [t for t in _REQUIRED_SPECIALS if t not in vocab] + if missing: + raise ValueError( + f"Tokenizer at {model_path} is missing required Higgs TTS v3 " + f"special tokens: {missing}. Available added tokens: " + f"{sorted(vocab.keys())[:20]}" + ) + + self.tts_token_id = vocab["<|tts|>"] + self.text_token_id = vocab["<|text|>"] + self.audio_continuation_id = vocab["<|audio|>"] + + if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None: + self.eos_token_id = int(tokenizer.eos_token_id) + + +AutoConfig.register("higgs_multimodal_qwen3", HiggsAudioV3Config) diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py new file mode 100644 index 000000000..40936f85d --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py @@ -0,0 +1,525 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage 1 (codec decoder) for higgs-audio v3. + +Reuses higgs-audio-v2's RVQ + DAC codec decoder but loads weights from the +v3 checkpoint's bundled codec (``tied.embedding.modality_embeddings.0.model.*`` +prefix) rather than from a standalone tokenizer repo. +""" + +from __future__ import annotations + +import os +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import ( + HiggsAudioRVQ, + build_boson_dac_decoder, + build_higgs_audio_acoustic_decoder, + load_higgs_audio_codec, +) +from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +__all__ = [ + "HiggsAudioV3Code2Wav", + "HiggsAudioV3Code2WavForConditionalGeneration", +] + +logger = init_logger(__name__) + + +# Prefix under which codec weights are stored in the v3 checkpoint. +_CODEC_PREFIX = "tied.embedding.modality_embeddings.0.model." + + +class HiggsAudioV3Code2Wav(nn.Module): + """Stage-1 codec decoder for higgs-audio v3. + + The codec architecture is identical to v2 (HiggsAudioRVQ + fc2 + + BosonDacDecoder); only the weight loading path differs. V3 bundles + codec weights inside the main checkpoint under a known prefix, while + v2 uses a standalone tokenizer repo. + + For backward compatibility and simplicity, we delegate to the v2 + ``load_higgs_audio_codec`` function when the audio tokenizer repo + is available. When not (or as primary path), we extract codec weights + from the v3 safetensors via the engine's weight iterator. + """ + + input_modalities = "audio" + + def __init__( + self, + config: HiggsAudioV3Config | None = None, + *, + vllm_config: VllmConfig | None = None, + prefix: str = "", + ): + super().__init__() + + if vllm_config is not None: + hf_config = vllm_config.model_config.hf_config + if isinstance(hf_config, HiggsAudioV3Config): + self.config = hf_config + else: + self.config = HiggsAudioV3Config(**hf_config.to_dict()) + self._model_path: str | None = vllm_config.model_config.model + self.vllm_config: VllmConfig | None = vllm_config + else: + if config is None: + raise TypeError("HiggsAudioV3Code2Wav: provide either `config` or `vllm_config`.") + self.config = config + self._model_path = None + self.vllm_config = None + + self.sample_rate: int = int(self.config.sample_rate) + self.num_codebooks: int = int(self.config.num_codebooks) + self.num_real_codes: int = int(self.config.num_real_codes) + self.hop_length: int = 960 + + # Engine-runner hooks + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + # Codec modules (populated by load_weights) + self.quantizer: HiggsAudioRVQ | None = None + self.fc2: nn.Linear | None = None + self.acoustic_decoder: nn.Module | None = None + self._loaded: bool = False + # Do NOT eagerly load from standalone tokenizer here. + # load_weights() will try bundled V3 codec first, then fall back. + + def _load_from_tokenizer_repo(self) -> None: + """Load codec from the standalone higgs-audio-v2-tokenizer repo.""" + tokenizer_id = "bosonai/higgs-audio-v2-tokenizer" + # Try to resolve from HF cache + from huggingface_hub import try_to_load_from_cache + + cached = try_to_load_from_cache(repo_id=tokenizer_id, filename="config.json") + if isinstance(cached, str) and os.path.isfile(cached): + tokenizer_dir = os.path.dirname(cached) + else: + from huggingface_hub.constants import HF_HUB_CACHE + + safe = tokenizer_id.replace("/", "--") + snapshots_dir = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots") + if os.path.isdir(snapshots_dir): + tokenizer_dir = None + for rev in os.listdir(snapshots_dir): + candidate = os.path.join(snapshots_dir, rev) + if os.path.isfile(os.path.join(candidate, "config.json")): + tokenizer_dir = candidate + break + if tokenizer_dir is None: + raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}") + else: + raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + quantizer, fc2, acoustic_decoder, _cfg = load_higgs_audio_codec(tokenizer_dir, device) + self.quantizer = quantizer + self.fc2 = fc2 + self.acoustic_decoder = acoustic_decoder + self._loaded = True + logger.info("Loaded HiggsAudioV3Code2Wav from standalone tokenizer repo.") + + # ------------------------------------------------------------------ engine hooks + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None: + return None + + # ------------------------------------------------------------------ load + def load_weights(self, weights_or_model_dir, device: torch.device | None = None): + if self._loaded: + if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)): + for _ in weights_or_model_dir: + pass + return {name for name, _ in self.named_parameters()} + + if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)): + # Engine path: try to extract V3 bundled codec weights first + codec_state: dict[str, torch.Tensor] = {} + for name, tensor in weights_or_model_dir: + if name.startswith(_CODEC_PREFIX): + stripped = name[len(_CODEC_PREFIX) :] + codec_state[stripped] = tensor + + if codec_state: + # Bundled keys found — load strictly, no fallback on error. + self._load_from_bundled_state(codec_state, device) + logger.info( + "Loaded HiggsAudioV3Code2Wav from bundled V3 checkpoint (%d codec keys consumed).", + len(codec_state), + ) + return {name for name, _ in self.named_parameters()} + + # No bundled keys — fallback to standalone tokenizer repo + if not self._loaded: + try: + self._load_from_tokenizer_repo() + except (FileNotFoundError, OSError) as exc2: + raise RuntimeError( + "HiggsAudioV3Code2Wav: could not load codec weights. " + "Neither V3 bundled codec keys (tied.embedding." + "modality_embeddings.0.model.*) were found in the " + "checkpoint, nor is bosonai/higgs-audio-v2-tokenizer " + f"cached in HF_HOME. Last error: {exc2}" + ) from exc2 + return {name for name, _ in self.named_parameters()} + + # File-based path + if not self._loaded: + self._load_from_tokenizer_repo() + return None + + def _load_from_bundled_state( + self, + codec_state: dict[str, torch.Tensor], + device: torch.device | None = None, + ) -> None: + """Load codec from V3 checkpoint's bundled weights. + + Keys arrive with the ``tied.embedding.modality_embeddings.0.model.`` + prefix already stripped. We remap to the same layout that + ``load_higgs_audio_codec`` expects and build the modules. + """ + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import ( + _remap_boson_model_pth_state_dict, + ) + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # The bundled codec may use boson-ai key naming. Try remapping. + needs_remap = any(k.startswith("quantizer.vq.layers.") for k in codec_state) + if needs_remap: + codec_state = _remap_boson_model_pth_state_dict(codec_state) + + # Discover parameters from the state dict + proj_key = "quantizer.quantizers.0.project_out.weight" + if proj_key not in codec_state: + raise KeyError( + f"Bundled codec state dict missing '{proj_key}'. " + f"Available keys (first 10): {list(codec_state.keys())[:10]}" + ) + hidden_size = codec_state[proj_key].shape[0] + codebook_dim = 64 + codebook_size = 1024 + required_quantizers = self.num_codebooks + + # Build RVQ — require exactly num_codebooks quantizers + quantizer = HiggsAudioRVQ( + num_quantizers=required_quantizers, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + hidden_size=hidden_size, + ).to(device) + for i in range(required_quantizers): + prefix = f"quantizer.quantizers.{i}" + # Require codebook weight + embed_key = None + for key_suffix in [".codebook.embed", ".codebook.weight"]: + candidate = f"{prefix}{key_suffix}" + if candidate in codec_state: + embed_key = candidate + break + if embed_key is None: + raise KeyError( + f"Bundled codec missing codebook weight for quantizer {i}. " + f"Expected {prefix}.codebook.embed or {prefix}.codebook.weight" + ) + quantizer.quantizers[i].codebook.weight.data.copy_(codec_state[embed_key].to(device)) + # Require projector weights + proj_w = f"{prefix}.project_out.weight" + proj_b = f"{prefix}.project_out.bias" + if proj_w not in codec_state: + raise KeyError(f"Bundled codec missing {proj_w}") + if proj_b not in codec_state: + raise KeyError(f"Bundled codec missing {proj_b}") + quantizer.quantizers[i].project_out.weight.data.copy_(codec_state[proj_w].to(device)) + quantizer.quantizers[i].project_out.bias.data.copy_(codec_state[proj_b].to(device)) + + # Build fc2 + fc2_w = codec_state.get("fc2.weight") + fc2_b = codec_state.get("fc2.bias") + if fc2_w is None or fc2_b is None: + raise KeyError( + "Bundled codec state dict missing fc2.weight and/or fc2.bias. " + f"Available keys (first 10): {list(codec_state.keys())[:10]}" + ) + fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device) + fc2.weight.data.copy_(fc2_w.to(device)) + fc2.bias.data.copy_(fc2_b.to(device)) + + # Build acoustic decoder — detect layout from key names. + # V3 checkpoint uses OmniVoice layout (block.N/conv1/conv2/snake1), + # not boson-ai standalone layout (model.0/model.1). + ad_keys = { + k[len("acoustic_decoder.") :]: v for k, v in codec_state.items() if k.startswith("acoustic_decoder.") + } + is_boson_layout = any(k.startswith("model.") for k in ad_keys) + if is_boson_layout: + acoustic_decoder = build_boson_dac_decoder(device) + else: + # OmniVoice layout — need the v2 tokenizer config for DacConfig + tokenizer_config = { + "acoustic_model_config": { + "codebook_dim": 8, + "codebook_size": 1024, + "decoder_hidden_size": 1024, + "downsampling_ratios": [8, 5, 4, 2, 3], + "encoder_hidden_size": 64, + "hidden_size": 256, + "hop_length": 960, + "model_type": "dac", + "n_codebooks": 9, + "sampling_rate": 16000, + "upsampling_ratios": [8, 5, 4, 2, 3], + } + } + acoustic_decoder = build_higgs_audio_acoustic_decoder(tokenizer_config, device) + if not ad_keys: + raise KeyError( + "Bundled codec state dict has no acoustic_decoder.* keys. " + f"Available keys (first 10): {list(codec_state.keys())[:10]}" + ) + load_report = acoustic_decoder.load_state_dict(ad_keys, strict=False) + if load_report.missing_keys: + raise RuntimeError( + f"Bundled DAC decoder missing {len(load_report.missing_keys)} keys: {load_report.missing_keys[:5]}..." + ) + if load_report.unexpected_keys: + raise RuntimeError( + f"Bundled DAC decoder has {len(load_report.unexpected_keys)} unexpected keys: " + f"{load_report.unexpected_keys[:5]}..." + ) + acoustic_decoder.eval() + + # Verify required decoder-side codec keys were consumed. The V3 + # checkpoint also bundles encoder-side keys (acoustic_encoder.*, + # encoder_semantic.*, decoder_semantic.*, semantic_model.*, fc.*, fc1.*) + # that the DAC decoder does not need — those are allowed but ignored. + consumed: set[str] = set() + for i in range(required_quantizers): + prefix = f"quantizer.quantizers.{i}" + for suffix in [".codebook.embed", ".codebook.weight"]: + if f"{prefix}{suffix}" in codec_state: + consumed.add(f"{prefix}{suffix}") + consumed.add(f"{prefix}.project_out.weight") + consumed.add(f"{prefix}.project_out.bias") + consumed.add("fc2.weight") + consumed.add("fc2.bias") + for k in codec_state: + if k.startswith("acoustic_decoder."): + consumed.add(k) + + # Known unused prefixes: encoder-side modules and quantizer training + # artifacts (cluster_size, embed_avg, inited, project_in) that are + # bundled in the checkpoint but not needed for inference decoding. + _KNOWN_UNUSED_PREFIXES = ( + "acoustic_encoder.", + "encoder_semantic.", + "decoder_semantic.", + "semantic_model.", + "fc.", + "fc1.", + ) + # Quantizer training-only suffixes (per-codebook) + _QUANTIZER_TRAINING_SUFFIXES = ( + ".codebook.cluster_size", + ".codebook.embed_avg", + ".codebook.inited", + ".project_in.weight", + ".project_in.bias", + ) + unconsumed = set() + for k in codec_state: + if k in consumed: + continue + if any(k.startswith(p) for p in _KNOWN_UNUSED_PREFIXES): + continue + if any(k.endswith(s) for s in _QUANTIZER_TRAINING_SUFFIXES): + continue + unconsumed.add(k) + if unconsumed: + raise RuntimeError( + f"Bundled codec has {len(unconsumed)} unexpected decoder-side keys: {sorted(unconsumed)[:10]}..." + ) + + self.quantizer = quantizer + self.fc2 = fc2 + self.acoustic_decoder = acoustic_decoder + self._loaded = True + + # ------------------------------------------------------------------ decode + @torch.inference_mode() + def decode_codes(self, audio_codes: torch.Tensor) -> torch.Tensor: + """Decode [B, num_codebooks=8, T] codes to PCM [B, 1, T*960].""" + if not self._loaded: + raise RuntimeError("HiggsAudioV3Code2Wav not loaded.") + + codes = self._validate_codes(audio_codes) + rvq_codes = codes.transpose(0, 1).long() + quantized = self.quantizer.decode(rvq_codes) + quantized = quantized.to(dtype=self.fc2.weight.dtype) + quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2) + first_param = next(self.acoustic_decoder.parameters(), None) + if first_param is not None and quantized.dtype != first_param.dtype: + quantized = quantized.to(dtype=first_param.dtype) + audio = self.acoustic_decoder(quantized) + if audio.dim() == 2: + audio = audio.unsqueeze(1) + return audio.to(codes.device) + + @torch.inference_mode() + def forward_chunk( + self, + audio_codes: torch.Tensor, + *, + left_context_size: int = 0, + right_holdback_size: int = 0, + hop_length: int | None = None, + ) -> torch.Tensor: + # ``left_context_size`` frames trimmed off the front (already-emitted + # prefix). ``right_holdback_size`` frames trimmed off the end + # (deferred to the next chunk so the codec has future context for + # the emit region's trailing samples). + hop = int(hop_length) if hop_length is not None else self.hop_length + pcm = self.decode_codes(audio_codes) + right_trim = right_holdback_size * hop + left_trim = left_context_size * hop + if right_trim > 0: + if pcm.shape[-1] <= right_trim: + return pcm[..., :0] + pcm = pcm[..., :-right_trim] + if left_trim == 0: + return pcm + if pcm.shape[-1] <= left_trim: + return pcm[..., :0] + return pcm[..., left_trim:] + + # ------------------------------------------------------------------ runtime forward + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> torch.Tensor | OmniOutput: + if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 3: + return self.decode_codes(input_ids) + + sr_val = int(self.sample_rate) + sr_tensor = torch.tensor(sr_val, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + + if input_ids is None or input_ids.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]}, + ) + + ids = input_ids.reshape(-1).to(dtype=torch.long) + request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) + + left_context_size = [0] * len(request_ids_list) + right_holdback_size = [0] * len(request_ids_list) + if runtime_additional_information is not None: + for i, info in enumerate(runtime_additional_information): + if i >= len(left_context_size): + break + meta = info.get("meta", {}) if isinstance(info, dict) else {} + if "left_context_size" in meta: + left_context_size[i] = int(meta["left_context_size"]) + if "right_holdback_size" in meta: + right_holdback_size[i] = int(meta["right_holdback_size"]) + + wavs: list[torch.Tensor] = [] + for i, req_ids in enumerate(request_ids_list): + n = int(req_ids.numel()) + if n == 0: + wavs.append(empty) + continue + if n % self.num_codebooks != 0: + logger.warning( + "HiggsAudioV3Code2Wav: flat code length %d not divisible by %d", + n, + self.num_codebooks, + ) + wavs.append(empty) + continue + frames = n // self.num_codebooks + codes_qf = req_ids.reshape(self.num_codebooks, frames) + codes_bqf = codes_qf.unsqueeze(0) + try: + pcm = self.forward_chunk( + codes_bqf, + left_context_size=left_context_size[i], + right_holdback_size=right_holdback_size[i], + hop_length=self.hop_length, + ) + except ValueError as exc: + logger.warning("HiggsAudioV3Code2Wav: decode skipped (%s)", exc) + wavs.append(empty) + continue + wavs.append(pcm.squeeze(0).squeeze(0).to(torch.float32).cpu()) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": wavs, + "sr": [sr_tensor] * len(wavs), + }, + ) + + # ------------------------------------------------------------------ helpers + def _validate_codes(self, audio_codes: torch.Tensor) -> torch.Tensor: + if audio_codes.ndim != 3: + raise ValueError( + f"audio_codes must have shape [B, {self.num_codebooks}, T]; got {tuple(audio_codes.shape)}" + ) + if int(audio_codes.shape[1]) != self.num_codebooks: + raise ValueError(f"dim 1 must equal num_codebooks={self.num_codebooks}; got {int(audio_codes.shape[1])}") + # ``.item()`` triggers a GPU->CPU sync and is illegal inside CUDA-graph + # capture (the dummy run before warmup hits this path). Skip the + # range check while a graph is being captured; the per-step validator + # still fires at real inference time. + if audio_codes.numel() > 0 and not torch.cuda.is_current_stream_capturing(): + max_val = int(audio_codes.max().item()) + min_val = int(audio_codes.min().item()) + if max_val >= self.num_real_codes or min_val < 0: + raise ValueError( + f"audio_codes out of range: min={min_val}, max={max_val}; expected [0, {self.num_real_codes - 1}]" + ) + return audio_codes + + @staticmethod + def _split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]: + if seq_token_counts is not None and len(seq_token_counts) > 1: + boundaries = [0] + for count in seq_token_counts: + boundaries.append(boundaries[-1] + int(count)) + n = int(ids.numel()) + return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))] + return [ids] + + +HiggsAudioV3Code2WavForConditionalGeneration = HiggsAudioV3Code2Wav diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py new file mode 100644 index 000000000..9d4e05356 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py @@ -0,0 +1,1524 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-0 talker for higgs-audio v3 (Qwen3 backbone, fused multi-codebook). + +Architecture: +- Backbone: Qwen3 (~4B, 36 layers, 2560 hidden, GQA 32/8). No DualFFN. +- Fused multi-codebook embedding: [N*V, D] weight, offset lookup, sum across N +- Fused multi-codebook head: same weight (tied), reshape to [L, N, V] +- MusicGen-style delay pattern [0,1,...,7] with BOC/EOC +- Audio feedback: replace continuation-token embedding with fused codebook embed + +Weight loading maps from the HF checkpoint's prefixes: + tied.embedding.text_embedding. -> model.embed_tokens. + body.layers. -> model.layers. + body.norm. -> model.norm. + tied.head.text_head. -> lm_head. + tied.embedding.modality_embeddings.0.embedding. -> multimodal_embedding. + tied.embedding.modality_embeddings.0.model.* -> skipped (codec for code2wav) + tied.head.modality_heads.0.* -> skipped when tied +""" + +from __future__ import annotations + +import copy +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.platforms import current_platform +from vllm.v1.outputs import SamplerOutput + +from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + HiggsAudioV3Config, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +__all__ = ["HiggsAudioV3TalkerForConditionalGeneration"] + +logger = init_logger(__name__) + + +# Delay pattern constants +BOC_ID = 1024 # beginning of codebook +EOC_ID = 1025 # end of codebook + +# Checkpoint prefix mapping +_BACKBONE_PREFIX_MAP = { + "tied.embedding.text_embedding.": "model.embed_tokens.", + "body.layers.": "model.layers.", + "body.norm.": "model.norm.", + "tied.head.text_head.": "lm_head.", +} +_MODALITY_EMBEDDING_PREFIX = "tied.embedding.modality_embeddings.0.embedding." +_MODALITY_HEAD_PREFIX = "tied.head.modality_heads.0." +_CODEC_PREFIX = "tied.embedding.modality_embeddings.0.model." + + +class HiggsFusedMultiTextEmbedding(nn.Module): + """Fused multi-codebook embedding: [N*V, D] weight + offset lookup.""" + + def __init__(self, num_codebooks: int, vocab_size: int, hidden_size: int): + super().__init__() + self.weight = nn.Parameter(torch.empty(num_codebooks * vocab_size, hidden_size)) + self.num_codebooks = num_codebooks + self.vocab_size = vocab_size + self.register_buffer( + "_offsets", + torch.arange(num_codebooks, dtype=torch.long) * vocab_size, + persistent=False, + ) + + def forward(self, codes: torch.Tensor) -> torch.Tensor: + fused_ids = codes + self._offsets + return F.embedding(fused_ids, self.weight).sum(dim=-2) + + +class HiggsFusedMultiTextHead(nn.Module): + """Fused multi-codebook head: [L, D] -> [L, N, V] via one linear.""" + + def __init__(self, num_codebooks: int, vocab_size: int, hidden_size: int): + super().__init__() + self.weight = nn.Parameter(torch.empty(num_codebooks * vocab_size, hidden_size)) + self.num_codebooks = num_codebooks + self.vocab_size = vocab_size + + def generate(self, hidden: torch.Tensor) -> torch.Tensor: + logits = F.linear(hidden, self.weight) + return logits.reshape(hidden.shape[0], self.num_codebooks, self.vocab_size) + + +class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): + """Stage-0 talker for higgs-audio v3. + + Wraps Qwen3Model backbone + fused multi-codebook modules for TTS generation + with MusicGen-style delay pattern sampling and audio feedback embedding. + """ + + # Tell the AR runner to call model.sample() instead of the stock sampler. + prefer_model_sampler: bool = True + # Tell the runner to call postprocess() to emit per-step audio codes. + have_multimodal_outputs: bool = True + has_postprocess: bool = True + skips_model_sampler_output_token_history: bool = True + postprocess_uses_hidden_states: bool = False + postprocess_uses_multimodal_outputs: bool = False + postprocess_uses_req_infos: bool = False + supports_omni_query_start_loc: bool = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + hf_config = vllm_config.model_config.hf_config + if isinstance(hf_config, HiggsAudioV3Config): + self.config = hf_config + else: + self.config = HiggsAudioV3Config(**hf_config.to_dict()) + model_config = getattr(vllm_config, "model_config", None) + compilation_config = getattr(vllm_config, "compilation_config", None) + cudagraph_mode = getattr(compilation_config, "cudagraph_mode", None) + mode_name = getattr(cudagraph_mode, "name", str(cudagraph_mode)).upper() + self._use_external_decode_cudagraph = ( + not bool(getattr(model_config, "enforce_eager", True)) + and cudagraph_mode is not None + and "NONE" not in mode_name + ) + if self._use_external_decode_cudagraph: + self.config.enable_mlp_cudagraph = False + assert not (self._use_external_decode_cudagraph and self.config.enable_mlp_cudagraph), ( + "Higgs Audio v3 local MLP graph and external decode CUDA graph are mutually exclusive" + ) + + self.vllm_config = vllm_config + self.num_codebooks = int(self.config.num_codebooks) + self.codebook_size = int(self.config.codebook_size) + hidden_size = int(self.config.audio_hidden_size) + self.tie_modality = self.config.tie_modality_embeddings + + # Fused multi-codebook modules + self.multimodal_embedding = HiggsFusedMultiTextEmbedding(self.num_codebooks, self.codebook_size, hidden_size) + self.modality_head = HiggsFusedMultiTextHead(self.num_codebooks, self.codebook_size, hidden_size) + if self.tie_modality: + self.modality_head.weight = self.multimodal_embedding.weight + + # Qwen3 backbone + self._backbone_config = self.config.text_config + backbone_vllm_config = copy.copy(vllm_config) + backbone_model_config = copy.copy(vllm_config.model_config) + backbone_model_config.hf_config = self._backbone_config + backbone_vllm_config.model_config = backbone_model_config + + self.model = Qwen3Model( + vllm_config=backbone_vllm_config, + prefix=f"{prefix}.model" if prefix else "model", + ) + if self.config.enable_flashinfer_api_unwrap: + self._maybe_unwrap_flashinfer_api_wrappers() + + if self._backbone_config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + self._backbone_config.vocab_size, + self._backbone_config.hidden_size, + prefix=f"{prefix}.lm_head" if prefix else "lm_head", + ) + + self.logits_processor = LogitsProcessor(self._backbone_config.vocab_size) + self._text_vocab_size = int(self._backbone_config.vocab_size) + + # Audio continuation token ID — resolved lazily from tokenizer. + # This is the <|audio|> token that serves as the LM-level continuation + # marker during audio generation (equivalent to v2's audio_token_id). + self._audio_continuation_id: int | None = None + self._eos_token_id: int | None = None + self._resolved_tokens = False + + self._last_logits_hidden: torch.Tensor | None = None + self._last_step_input_ids: torch.Tensor | None = None + self._last_step_query_start_loc: torch.Tensor | None = None + self._last_step_query_start_loc_buffer: torch.Tensor | None = None + self.supports_omni_decode_step_metadata = True + self._decode_step_metadata_from_runner = False + self._last_audio_codes: torch.Tensor | None = None + self._last_audio_code_valid: list[bool] = [] + self._postprocess_cursor: int = 0 + self._last_audio_codes_buffer: torch.Tensor | None = None + self._last_audio_host_staging: torch.Tensor | None = None + self._last_audio_gpu_staging: torch.Tensor | None = None + self._last_audio_staging_event: torch.cuda.Event | None = None + self._last_audio_valid_flags: list[int] | None = None + self._last_audio_done_flags: list[int] | None = None + self._audio_staging_events: list[torch.cuda.Event] = [] + self._audio_staging_event_cursor: int = 0 + self._row_index_cache: dict[tuple[str, int], torch.Tensor] = {} + self._codebook_index_cache: dict[tuple[str, int], torch.Tensor] = {} + self._boc_frame_cache: dict[tuple[str, int], torch.Tensor] = {} + self._fast_audio_direct_rows: int = 0 + scheduler_config = getattr(vllm_config, "scheduler_config", None) + self._mlp_cudagraph_max_batch = max(1, int(getattr(scheduler_config, "max_num_seqs", 16))) + self._postprocess_audio_rows: int = 0 + self._postprocess_audio_active_rows: int = 0 + self._mlp_graphs: dict[tuple[int, int], dict[str, Any]] = {} + self._mlp_graph_disabled: set[tuple[int, int]] = set() + + # Pre-allocated decode-step audio feedback buffers (CUDA-graph safe). + # Populated by sample(), read by forward() via torch.where (no dict). + max_bs = max(64, self._mlp_cudagraph_max_batch) + self.register_buffer( + "_decode_last_codes", + torch.zeros(max_bs, self.num_codebooks, dtype=torch.long), + persistent=False, + ) + self.register_buffer("_decode_has_codes", torch.zeros(max_bs, dtype=torch.bool), persistent=False) + # Keep delay/ramp counters as int64 to match arange/codebook-index + # arithmetic in the decode hot path without per-step dtype casts. + self.register_buffer("_decode_delay_count", torch.zeros(max_bs, dtype=torch.long), persistent=False) + self.register_buffer("_decode_eoc_countdown", torch.full((max_bs,), -1, dtype=torch.long), persistent=False) + self.register_buffer("_decode_generation_done", torch.zeros(max_bs, dtype=torch.bool), persistent=False) + self._decode_active_audio_count: int = 0 + + # PrefixCache opt-outs (mirror qwen3_tts pattern): + # 1. The talker only consumes the last token's hidden state, so the + # runner can skip the per-step full hidden-state GPU->CPU merge + # that PrefixCache otherwise does. + # 2. Per-step ``codes.audio`` rows stay GPU-resident; defer the CPU + # write of the prefix-cache mm-output copy to request finish so + # the per-step bookkeeping does not block batching. Stage 0 can + # then set ``enable_prefix_caching: true`` without the regression + # observed in qwen3_tts (#3665). + self.requires_full_prefix_cached_hidden_states = False + self.deferred_prefix_cache_mm_keys = {"codes.audio"} + + def _maybe_unwrap_flashinfer_api_wrappers(self) -> None: + """Remove FlashInfer trace wrappers only from wrappers owned by this model.""" + patched: list[str] = [] + saw_flashinfer = False + seen: set[int] = set() + stack: list[Any] = [] + for layer in getattr(self.model, "layers", []): + impl = getattr(getattr(getattr(layer, "self_attn", None), "attn", None), "impl", None) + if impl is not None: + stack.append(impl) + + while stack and len(seen) < 512: + obj = stack.pop() + obj_id = id(obj) + if obj_id in seen: + continue + seen.add(obj_id) + + cls = obj.__class__ + if "flashinfer" in cls.__module__.lower(): + saw_flashinfer = True + for method_name in ("run", "paged_run", "plan"): + method = getattr(obj, method_name, None) + original = getattr(method, "__wrapped__", None) + if original is None: + continue + try: + bound = original.__get__(obj, cls) if hasattr(original, "__get__") else original + setattr(obj, method_name, bound) + except Exception: + continue + patched.append(f"{cls.__name__}.{method_name}") + + obj_dict = getattr(obj, "__dict__", None) + if not isinstance(obj_dict, dict): + continue + for value in obj_dict.values(): + if isinstance(value, dict): + stack.extend(value.values()) + elif isinstance(value, (list, tuple, set)): + stack.extend(value) + elif not isinstance(value, (bool, int, float, str, bytes, type(None), torch.Tensor)): + stack.append(value) + + if patched: + logger.info("HiggsAudioV3Talker: unwrapped FlashInfer API wrappers: %s", sorted(set(patched))) + elif saw_flashinfer: + logger.warning( + "HiggsAudioV3Talker: FlashInfer backend was detected but no API wrappers were unwrapped; " + "FlashInfer internals may have changed and wrapper-unwrapping performance benefits may be unavailable." + ) + + def _resolve_token_ids(self) -> None: + """Resolve <|audio|> and eos token IDs. + + Prefers config's pre-resolved IDs (from ``resolve_special_tokens()``), + falls back to loading the HF tokenizer directly. + """ + if self._resolved_tokens: + return + self._resolved_tokens = True + + # Try config first (populated by resolve_special_tokens or from_pretrained) + cfg_audio = getattr(self.config, "audio_continuation_id", None) + cfg_eos = getattr(self.config, "eos_token_id", None) + if cfg_audio is not None: + self._audio_continuation_id = int(cfg_audio) + if cfg_eos is not None: + self._eos_token_id = int(cfg_eos) + + if self._audio_continuation_id is not None: + logger.info( + "Resolved v3 token IDs from config: audio_continuation=%s, eos=%s", + self._audio_continuation_id, + self._eos_token_id, + ) + return + + # Fallback: load tokenizer directly + model_path = getattr(self.vllm_config.model_config, "model", None) + if model_path is None: + return + try: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + vocab = dict(tokenizer.get_added_vocab()) + if "<|audio|>" in vocab: + self._audio_continuation_id = vocab["<|audio|>"] + if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None: + self._eos_token_id = int(tokenizer.eos_token_id) + logger.info( + "Resolved v3 token IDs from tokenizer: audio_continuation=%s, eos=%s", + self._audio_continuation_id, + self._eos_token_id, + ) + except Exception as exc: + logger.warning("Failed to resolve token IDs from tokenizer: %s", exc) + + def update_decode_step_metadata( + self, + *, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + omni_query_start_loc: torch.Tensor | None = None, + **_: Any, + ) -> None: + """Update per-step metadata before runner forward or CUDA graph replay.""" + _ = positions, inputs_embeds + if input_ids is not None: + self._last_step_input_ids = input_ids + if self._use_external_decode_cudagraph and input_ids.is_cuda and input_ids.ndim >= 1: + self._ensure_decode_state_capacity(int(input_ids.shape[0]), input_ids.device) + self._set_last_step_query_start_loc(omni_query_start_loc) + self._decode_step_metadata_from_runner = True + + # ------------------------------------------------------------------ forward + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + metadata_from_runner = self._decode_step_metadata_from_runner + self._decode_step_metadata_from_runner = False + + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") + + if inputs_embeds is None: + # Mask -100 placeholders to 0 before embedding. Use torch.where + # (no Python data-dependent branch) so this is CUDA-graph safe. + safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids) + hidden_states = self.model.embed_tokens(safe_ids) + else: + hidden_states = inputs_embeds + + if input_ids is not None and not metadata_from_runner: + self._last_step_input_ids = input_ids + + # Stash query_start_loc for audio-state row mapping and max_query_len + # for prefill detection. Some attention backends do not expose + # query_start_loc, so prefer backend metadata when available and + # otherwise use the runner-supplied buffer. + max_query_len = None + if not metadata_from_runner: + try: + fallback_qsl = kwargs.get("omni_query_start_loc") + from vllm.forward_context import get_forward_context + + attn_metadata = get_forward_context().attn_metadata + if isinstance(attn_metadata, dict) and attn_metadata: + attn = next(iter(attn_metadata.values())) + else: + attn = attn_metadata + qsl = getattr(attn, "query_start_loc", None) + self._set_last_step_query_start_loc(qsl if isinstance(qsl, torch.Tensor) else fallback_qsl) + max_query_len = getattr(attn, "max_query_len", None) + except Exception: + self._last_step_query_start_loc = None + + # Prefill-only operations: ref audio substitution and audio feedback + # require Python dict/list ops that break CUDA graph capture. + # Prefer max_query_len where available: concurrent decode has + # input_ids.numel() > 1 but max_query_len == 1. + if max_query_len is not None: + is_prefill = int(max_query_len) > 1 + else: + is_prefill = ( + input_ids is not None + and inputs_embeds is None + and not self._is_single_token_decode_step(int(hidden_states.shape[0])) + ) + if is_prefill and info_dicts: + # Voice clone: replace -100 placeholder positions with ref audio embeddings + hidden_states = self._apply_ref_audio_substitution(hidden_states, input_ids, info_dicts) + + # Audio feedback at decode: replace continuation token embeddings + if input_ids is not None and inputs_embeds is None: + hidden_states = self._apply_audio_feedback(hidden_states, input_ids) + + if self.config.enable_mlp_cudagraph: + return self._run_qwen3_layers_mlp_graph(positions, hidden_states) + + return self._run_qwen3_layers_eager(positions, hidden_states) + + def _run_qwen3_layers_eager(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + residual: torch.Tensor | None = None + for layer in self.model.layers: + hidden_states, residual = layer(positions, hidden_states, residual) + norm_out = self.model.norm(hidden_states, residual) + if isinstance(norm_out, tuple): + norm_out = norm_out[0] + return norm_out + + def _run_qwen3_layers_mlp_graph(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + residual: torch.Tensor | None = None + batch = int(hidden_states.shape[0]) if hidden_states.ndim >= 1 else 0 + use_graph = ( + torch.cuda.is_available() + and hidden_states.is_cuda + and hidden_states.ndim == 2 + and positions.ndim == 1 + and batch > 0 + and batch <= self._mlp_cudagraph_max_batch + and int(positions.shape[0]) == batch + and self._is_decode_only_graph_batch(batch) + ) + for layer_idx, layer in enumerate(self.model.layers): + if residual is None: + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + else: + hidden_states, residual = layer.input_layernorm(hidden_states, residual) + hidden_states = layer.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + if use_graph: + graph_out = self._run_qwen3_mlp_graph(layer_idx, layer, hidden_states, residual) + if graph_out is not None: + hidden_states, residual = graph_out + continue + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp(hidden_states) + norm_out = self.model.norm(hidden_states, residual) + if isinstance(norm_out, tuple): + norm_out = norm_out[0] + return norm_out + + def _run_qwen3_mlp_graph( + self, + layer_idx: int, + layer: nn.Module, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + key = (int(layer_idx), int(hidden_states.shape[0])) + if key in self._mlp_graph_disabled: + return None + if hidden_states.ndim != 2 or residual.ndim != 2 or hidden_states.shape != residual.shape: + self._mlp_graph_disabled.add(key) + return None + if torch.cuda.is_current_stream_capturing(): + return None + + entry = self._mlp_graphs.get(key) + if entry is not None: + entry["static_hidden"].copy_(hidden_states) + entry["static_residual"].copy_(residual) + entry["graph"].replay() + return entry["static_output"], entry["static_residual_out"] + + try: + static_hidden = torch.empty_like(hidden_states) + static_residual = torch.empty_like(residual) + static_hidden.copy_(hidden_states) + static_residual.copy_(residual) + graph = torch.cuda.CUDAGraph() + with torch.inference_mode(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + normed, static_residual_out = layer.post_attention_layernorm(static_hidden, static_residual) + static_output = layer.mlp(normed) + self._mlp_graphs[key] = { + "graph": graph, + "static_hidden": static_hidden, + "static_residual": static_residual, + "static_output": static_output, + "static_residual_out": static_residual_out, + } + logger.info( + "HiggsAudioV3Talker: captured MLP CUDA graph for layer=%d decode batch=%d", + layer_idx, + int(hidden_states.shape[0]), + ) + return static_output, static_residual_out + except Exception as exc: + self._mlp_graph_disabled.add(key) + logger.warning( + "HiggsAudioV3Talker: MLP CUDA graph disabled for layer=%d batch=%d: %s", + layer_idx, + int(hidden_states.shape[0]), + exc, + ) + return None + + def _is_decode_only_graph_batch(self, batch: int) -> bool: + q_start = self._last_step_query_start_loc + return isinstance(q_start, torch.Tensor) and int(q_start.numel()) == batch + 1 + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: + self._last_logits_hidden = hidden_states + return self.logits_processor(self.lm_head, hidden_states, sampling_metadata) + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids) + text_embed = self.model.embed_tokens(safe_ids) + return self._apply_audio_feedback(text_embed, input_ids) + + def _set_last_step_query_start_loc(self, query_start_loc: Any) -> None: + if not isinstance(query_start_loc, torch.Tensor): + self._last_step_query_start_loc = None + return + source = query_start_loc.detach() + numel = int(source.numel()) + buf = self._last_step_query_start_loc_buffer + if buf is None or buf.device != source.device or buf.dtype != source.dtype or int(buf.numel()) < numel: + buf = torch.empty(max(numel, 17), dtype=source.dtype, device=source.device) + self._last_step_query_start_loc_buffer = buf + buf[:numel].copy_(source.reshape(-1)) + self._last_step_query_start_loc = buf[:numel] + + # ------------------------------------------------------------------ ref audio substitution + def _apply_ref_audio_substitution( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + info_dicts: list[dict[str, Any]] | None, + ) -> torch.Tensor: + """Replace -100 placeholder positions with fused multi-codebook embeddings + of the delay-pattern-encoded reference audio codes. + + Called at prefill to inject voice clone reference. ``info_dicts`` is a + list of per-request dicts from ``model_intermediate_buffer``, each + containing ``audio_input_ids`` ([T, N] delayed codes) and + ``audio_input_ids_mask`` ([T] bool mask). + """ + if not info_dicts: + return hidden_states + + PLACEHOLDER = -100 + flat_ids = input_ids.reshape(-1) + placeholder_mask = flat_ids == PLACEHOLDER + if not placeholder_mask.any(): + return hidden_states + + # Use query_start_loc to map placeholders to per-request spans + q_start = self._last_step_query_start_loc + if not isinstance(q_start, torch.Tensor) or q_start.numel() < 2: + # Fallback: single-request batch + q_start_list = [0, int(flat_ids.numel())] + else: + q_start_list = q_start.detach().to("cpu").tolist() + + new_hidden: torch.Tensor | None = None + num_requests = min(len(info_dicts), len(q_start_list) - 1) + + for i in range(num_requests): + info = info_dicts[i] + if not isinstance(info, dict): + continue + + codes = info.get("audio_input_ids") + mask = info.get("audio_input_ids_mask") + + # Handle msgspec serialization (may be list-wrapped) + if isinstance(codes, list): + codes = codes[0] if codes else None + if isinstance(mask, list): + mask = mask[0] if mask else None + if not isinstance(codes, torch.Tensor): + continue + + # codes shape: [T, num_codebooks] delayed reference codes + if codes.ndim == 3: + codes = codes[0] + if codes.ndim != 2: + continue + + if isinstance(mask, torch.Tensor): + if mask.ndim == 2: + mask = mask[0] + codes = codes[mask.to(dtype=torch.bool)] + + if codes.numel() == 0: + continue + + # Find placeholder positions in this request's span + s = int(q_start_list[i]) + e = int(q_start_list[i + 1]) + if e - s <= 1: + continue # Decode step, skip + + span_mask = placeholder_mask[s:e] + placeholders = span_mask.nonzero(as_tuple=True)[0] + n_codes = int(codes.shape[0]) + + if int(placeholders.numel()) < n_codes: + continue # Mismatch + + # Embed delayed codes via fused multi-codebook embedding + target = placeholders[:n_codes] + s + codes_device = codes.to(device=hidden_states.device, dtype=torch.long) + embeds = self.multimodal_embedding(codes_device) # [n_codes, hidden] + + if new_hidden is None: + new_hidden = hidden_states.clone() + flat_hidden = new_hidden.reshape(-1, new_hidden.shape[-1]) + flat_hidden[target] = embeds.to(new_hidden.dtype) + + return new_hidden if new_hidden is not None else hidden_states + + # ------------------------------------------------------------------ audio feedback + def _apply_audio_feedback(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + """Replace decode-step embeddings with audio feedback from pre-allocated buffers. + + CUDA-graph safe: reads from pre-allocated _decode_last_codes and + _decode_has_codes tensors using torch.where (no Python dict lookup). + The buffers are populated by sample() after each step. + + For decode steps (1 token per request): position i maps to row i. + For prefill: audio feedback is not needed (ref audio substitution + handles the prefill path separately). + """ + num_tokens = int(hidden_states.shape[0]) + external_decode_graph = bool(getattr(self, "_use_external_decode_cudagraph", False)) + if self._decode_active_audio_count == 0 and not external_decode_graph: + return hidden_states + + req_rows, token_positions = self._decode_request_token_positions(num_tokens, hidden_states.device) + if int(req_rows.numel()) == 0: + return hidden_states + + num_reqs = self._step_request_count(num_tokens) + self._ensure_decode_state_capacity(num_reqs, hidden_states.device) + + if self._is_single_token_decode_step(num_tokens): + codes_slice = self._decode_last_codes[:num_tokens] + has_codes = self._decode_has_codes[:num_tokens].unsqueeze(-1) + current_hidden = hidden_states + else: + codes_slice = self._decode_last_codes.index_select(0, req_rows) + has_codes = self._decode_has_codes.index_select(0, req_rows).unsqueeze(-1) + current_hidden = hidden_states.index_select(0, token_positions) + audio_embeds = self.multimodal_embedding(codes_slice) # [bs, D] + audio_embeds = audio_embeds.to(dtype=hidden_states.dtype) + replaced = torch.where(has_codes, audio_embeds, current_hidden) + + if self._is_single_token_decode_step(num_tokens): + return replaced + new_hidden = hidden_states.clone() + new_hidden.index_copy_(0, token_positions, replaced) + return new_hidden + + def _is_single_token_decode_step(self, num_rows: int) -> bool: + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor): + return int(q_start.numel()) == int(num_rows) + 1 + ids = getattr(self, "_last_step_input_ids", None) + if not isinstance(ids, torch.Tensor) or int(ids.numel()) != int(num_rows): + return False + if bool(getattr(self, "_use_external_decode_cudagraph", False)): + return True + return int(num_rows) == 1 + + def _step_request_count(self, num_tokens: int) -> int: + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) >= 2: + return int(q_start.numel()) - 1 + return int(num_tokens) + + def _decode_request_token_positions( + self, + num_tokens: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) >= 2: + q = q_start.to(device=device, dtype=torch.long) + starts = q[:-1] + spans = q[1:] - q[:-1] + req_rows = torch.arange(int(q.numel()) - 1, dtype=torch.long, device=device) + decode_mask = (spans == 1) & (starts >= 0) & (starts < int(num_tokens)) + return req_rows[decode_mask], starts[decode_mask] + + ids = getattr(self, "_last_step_input_ids", None) + if isinstance(ids, torch.Tensor) and int(ids.numel()) == int(num_tokens): + if bool(getattr(self, "_use_external_decode_cudagraph", False)): + rows = torch.arange(int(num_tokens), dtype=torch.long, device=device) + return rows, rows + if int(num_tokens) == 1: + row = torch.zeros(1, dtype=torch.long, device=device) + return row, row + empty = torch.empty(0, dtype=torch.long, device=device) + return empty, empty + + def _prefill_row_mask(self, num_rows: int, device: torch.device) -> torch.Tensor: + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) >= 2: + q = q_start.to(device=device, dtype=torch.long) + return (q[1:] - q[:-1]) > 1 + return torch.ones(num_rows, dtype=torch.bool, device=device) + + def _reset_decode_state_rows(self, mask: torch.Tensor, num_rows: int, device: torch.device) -> None: + self._ensure_decode_state_capacity(num_rows, device) + row_mask = mask.to(device=device, dtype=torch.bool) + self._decode_has_codes[:num_rows].masked_fill_(row_mask, False) + self._decode_generation_done[:num_rows].masked_fill_(row_mask, False) + self._decode_delay_count[:num_rows].masked_fill_(row_mask, 0) + self._decode_eoc_countdown[:num_rows].masked_fill_(row_mask, -1) + + def _ensure_decode_state_capacity(self, num_rows: int, device: torch.device | None = None) -> None: + """Keep GPU-resident decode state tensors aligned and large enough.""" + if device is not None and self._decode_last_codes.device != device: + self._decode_last_codes = self._decode_last_codes.to(device) + self._decode_has_codes = self._decode_has_codes.to(device) + self._decode_delay_count = self._decode_delay_count.to(device) + self._decode_eoc_countdown = self._decode_eoc_countdown.to(device) + self._decode_generation_done = self._decode_generation_done.to(device) + + cur_rows = int(self._decode_last_codes.shape[0]) + if num_rows <= cur_rows: + return + + new_size = max(num_rows, cur_rows * 2) + state_device = self._decode_last_codes.device + + new_last_codes = torch.zeros(new_size, self.num_codebooks, dtype=torch.long, device=state_device) + new_last_codes[:cur_rows].copy_(self._decode_last_codes) + self._decode_last_codes = new_last_codes + + new_has_codes = torch.zeros(new_size, dtype=torch.bool, device=state_device) + new_has_codes[:cur_rows].copy_(self._decode_has_codes) + self._decode_has_codes = new_has_codes + + new_delay_count = torch.zeros(new_size, dtype=torch.long, device=state_device) + new_delay_count[:cur_rows].copy_(self._decode_delay_count) + self._decode_delay_count = new_delay_count + + new_eoc_countdown = torch.full((new_size,), -1, dtype=torch.long, device=state_device) + new_eoc_countdown[:cur_rows].copy_(self._decode_eoc_countdown) + self._decode_eoc_countdown = new_eoc_countdown + + new_generation_done = torch.zeros(new_size, dtype=torch.bool, device=state_device) + new_generation_done[:cur_rows].copy_(self._decode_generation_done) + self._decode_generation_done = new_generation_done + + def _audio_seed_mask_from_step_input(self, num_rows: int, device: torch.device) -> torch.Tensor | None: + """Return rows whose previous token is <|audio|> using current step input_ids. + + This is the hot-path replacement for scanning sampling_metadata + output_token_ids. In decode, vLLM feeds the previous sampled token as + the current single-token input for each row. + """ + audio_id = self._audio_continuation_id + ids = self._last_step_input_ids + if audio_id is None or not isinstance(ids, torch.Tensor) or int(ids.numel()) <= 0: + return None + + flat_ids = ids.reshape(-1).to(device=device) + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) == num_rows + 1: + q = q_start.to(device=device, dtype=torch.long) + tail_idx = q[1:] - 1 + valid_tail = tail_idx >= 0 + tail_ids = flat_ids.index_select(0, tail_idx.clamp_min(0)) + return (tail_ids == int(audio_id)) & valid_tail + if int(flat_ids.numel()) == num_rows: + tail_ids = flat_ids + else: + return None + return tail_ids == int(audio_id) + + def _fast_audio_sampler_gpu_fallback_reason( + self, + *, + logits: torch.Tensor, + sampling_metadata: Any, + num_rows: int, + ) -> str | None: + if logits is None or logits.ndim != 2 or int(logits.shape[0]) != num_rows: + return "invalid_logits" + if num_rows <= 0: + return "empty_batch" + if getattr(sampling_metadata, "max_num_logprobs", None) is not None: + return "logprobs" + if getattr(sampling_metadata, "allowed_token_ids_mask", None) is not None: + return "allowed_token_ids" + if bool(getattr(sampling_metadata, "bad_words_token_ids", None)): + return "bad_words" + return None + + def _clear_last_audio_outputs(self) -> None: + self._last_audio_codes = None + self._last_audio_code_valid = [] + self._last_audio_host_staging = None + self._last_audio_staging_event = None + self._decode_active_audio_count = 0 + + def _apply_audio_mode_bias_batched( + self, + logits: torch.Tensor, + audio_mask: torch.Tensor, + done_mask: torch.Tensor, + ) -> None: + """Force audio/EOS rows with GPU masks, avoiding metadata scans.""" + audio_id = self._audio_continuation_id + if audio_id is None or logits is None or logits.ndim != 2: + return + + vocab = int(logits.shape[-1]) + if audio_id < 0 or audio_id >= vocab: + return + + target = torch.full( + (int(logits.shape[0]),), + int(audio_id), + dtype=torch.long, + device=logits.device, + ) + if self._eos_token_id is not None and 0 <= int(self._eos_token_id) < vocab: + eos_target = torch.full_like(target, int(self._eos_token_id)) + target = torch.where(done_mask.to(device=logits.device), eos_target, target) + + target = target.unsqueeze(1) + target_values = logits.gather(1, target) + logits.masked_fill_(audio_mask.to(device=logits.device).unsqueeze(1), float("-inf")) + logits.scatter_(1, target, target_values) + + # ------------------------------------------------------------------ sampling + def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> Any: + """Model-owned sampler with delay-pattern audio dispatch. + + Mirrors v2's pattern: bias LM logits to force audio continuation, + sample multi-codebook codes via the fused head, apply delay pattern, + and accumulate per-request state. + """ + self._resolve_token_ids() + + audio_id = self._audio_continuation_id + + def run_stock_sampler() -> Any: + sampler = getattr(self, "_stock_sampler", None) + if sampler is None: + from vllm.v1.sample.sampler import Sampler + + sampler = Sampler() + self._stock_sampler = sampler + return sampler(logits=logits, sampling_metadata=sampling_metadata) + + hidden = self._last_logits_hidden + self._last_logits_hidden = None + if hidden is None or audio_id is None: + sampler_output = run_stock_sampler() + self._clear_last_audio_outputs() + return sampler_output + + num_rows = int(hidden.shape[0]) + self._ensure_decode_state_capacity(num_rows, hidden.device) + ids = getattr(self, "_last_step_input_ids", None) + decode_only = isinstance(ids, torch.Tensor) and int(ids.numel()) == num_rows + if not decode_only: + self._fast_audio_direct_rows = 0 + pfmask = self._prefill_row_mask(num_rows, hidden.device) + self._reset_decode_state_rows(pfmask, num_rows, hidden.device) + prev_audio_mask = self._audio_seed_mask_from_step_input(num_rows, hidden.device) + if prev_audio_mask is None: + prev_audio_mask = torch.zeros(num_rows, dtype=torch.bool, device=hidden.device) + active_mask = self._decode_has_codes[:num_rows].to(device=hidden.device, dtype=torch.bool) + done_mask = self._decode_generation_done[:num_rows].to(device=hidden.device, dtype=torch.bool) + fast_fallback_reason = self._fast_audio_sampler_gpu_fallback_reason( + logits=logits, + sampling_metadata=sampling_metadata, + num_rows=num_rows, + ) + seed_mask = prev_audio_mask & ~active_mask & ~done_mask + audio_mask = prev_audio_mask | active_mask | done_mask + code_row_mask = (seed_mask | active_mask) & ~done_mask + gpu_stock_sampler_reasons = {"logprobs", "allowed_token_ids", "bad_words"} + use_gpu_audio_mode = fast_fallback_reason is None or fast_fallback_reason in gpu_stock_sampler_reasons + + if fast_fallback_reason is None: + direct_audio_batch = decode_only and self._fast_audio_direct_rows == num_rows + + if direct_audio_batch: + self._fast_audio_direct_rows = num_rows + audio_tokens = torch.full((num_rows,), int(audio_id), dtype=torch.int32, device=hidden.device) + if self._eos_token_id is not None: + eos_tokens = torch.full_like(audio_tokens, int(self._eos_token_id)) + audio_tokens = torch.where(done_mask, eos_tokens, audio_tokens) + sampled = audio_tokens.unsqueeze(-1) + sampler_output = SamplerOutput(sampled_token_ids=sampled, logprobs_tensors=None) + else: + self._apply_audio_mode_bias_batched(logits, audio_mask, done_mask) + sampler_output = run_stock_sampler() + sampled = getattr(sampler_output, "sampled_token_ids", None) + if sampled is None: + self._clear_last_audio_outputs() + return sampler_output + elif use_gpu_audio_mode: + self._fast_audio_direct_rows = 0 + self._apply_audio_mode_bias_batched(logits, audio_mask, done_mask) + sampler_output = run_stock_sampler() + sampled = getattr(sampler_output, "sampled_token_ids", None) + if sampled is None: + self._clear_last_audio_outputs() + return sampler_output + else: + self._fast_audio_direct_rows = 0 + sampler_output = run_stock_sampler() + self._clear_last_audio_outputs() + return sampler_output + + sampled_flat = sampled.reshape(-1) + if int(sampled_flat.numel()) != num_rows: + self._clear_last_audio_outputs() + return sampler_output + + if use_gpu_audio_mode: + audio_row_tensor = self._get_row_indices(num_rows, hidden.device) + seed_mask_1d = seed_mask.to(device=hidden.device, dtype=torch.bool) + done_mask_1d = done_mask.to(device=hidden.device, dtype=torch.bool) + self._decode_generation_done[:num_rows].masked_fill_(done_mask_1d, False) + self._decode_has_codes[:num_rows].masked_fill_(done_mask_1d, False) + self._decode_delay_count[:num_rows].masked_fill_(done_mask_1d, 0) + self._decode_eoc_countdown[:num_rows].masked_fill_(done_mask_1d, -1) + + boc_frames = self._get_boc_frames(num_rows, hidden.device) + seed_mask_2d = seed_mask_1d.unsqueeze(1) + self._decode_last_codes[:num_rows] = torch.where( + seed_mask_2d, + boc_frames, + self._decode_last_codes[:num_rows], + ) + self._decode_has_codes[:num_rows].masked_fill_(seed_mask_1d, True) + self._decode_delay_count[:num_rows].masked_fill_(seed_mask_1d, 0) + self._decode_eoc_countdown[:num_rows].masked_fill_(seed_mask_1d, -1) + self._decode_generation_done[:num_rows].masked_fill_(seed_mask_1d, False) + self._decode_active_audio_count = num_rows + else: + audio_row_tensor = self._get_row_indices(num_rows, hidden.device) + + # Per-codebook logits at audio positions + all_audio_rows = use_gpu_audio_mode + if use_gpu_audio_mode: + cb_logits = self._audio_codebook_logits_from_rows(hidden, audio_row_tensor, all_rows=True) + else: + cb_logits = self._audio_codebook_logits_from_rows(hidden, audio_row_tensor, all_rows=False) + + # Apply delay pattern masking BEFORE sampling + self._apply_delay_pattern_masking_batched(cb_logits, audio_row_tensor, all_rows=all_audio_rows) + + # Sample per-codebook + cb_logits_2d = cb_logits.reshape(-1, cb_logits.shape[-1]) + codes_2d = self._sample_audio_codes(cb_logits_2d) + codes_flat = codes_2d.view(cb_logits.shape[0], cb_logits.shape[1]).to(torch.long) + + self._update_delay_state_batched( + codes_flat, + audio_row_tensor, + num_rows, + hidden.device, + code_row_mask=code_row_mask, + all_rows=all_audio_rows, + ) + return sampler_output + + def _get_audio_codes_buffer(self, num_rows: int, device: torch.device) -> torch.Tensor: + buf = self._last_audio_codes_buffer + if buf is None or buf.device != device or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + buf = torch.empty((rows, self.num_codebooks), dtype=torch.long, device=device) + self._last_audio_codes_buffer = buf + return buf[:num_rows] + + def _get_audio_gpu_staging_buffer(self, num_rows: int, device: torch.device) -> torch.Tensor: + width = self.num_codebooks + 2 + buf = self._last_audio_gpu_staging + if buf is None or buf.device != device or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + buf = torch.empty((rows, width), dtype=torch.int32, device=device) + self._last_audio_gpu_staging = buf + return buf[:num_rows] + + def _get_audio_host_staging_buffer(self, num_rows: int) -> torch.Tensor: + width = self.num_codebooks + 2 + buf = self._last_audio_host_staging + if buf is None or buf.device.type != "cpu" or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + pin_memory = torch.cuda.is_available() + buf = torch.empty((rows, width), dtype=torch.int32, device="cpu", pin_memory=pin_memory) + self._last_audio_host_staging = buf + return buf[:num_rows] + + def _next_audio_staging_event(self) -> torch.cuda.Event: + if not self._audio_staging_events: + self._audio_staging_events = [torch.cuda.Event(), torch.cuda.Event()] + event = self._audio_staging_events[self._audio_staging_event_cursor] + self._audio_staging_event_cursor = (self._audio_staging_event_cursor + 1) % len(self._audio_staging_events) + return event + + @staticmethod + def _device_cache_key(device: torch.device) -> tuple[str, int]: + return (device.type, -1 if device.index is None else int(device.index)) + + def _get_row_indices(self, num_rows: int, device: torch.device) -> torch.Tensor: + key = self._device_cache_key(device) + buf = self._row_index_cache.get(key) + if buf is None or buf.device != device or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + buf = torch.arange(rows, dtype=torch.long, device=device) + self._row_index_cache[key] = buf + return buf[:num_rows] + + def _get_codebook_indices(self, device: torch.device) -> torch.Tensor: + key = self._device_cache_key(device) + buf = self._codebook_index_cache.get(key) + if buf is None or buf.device != device or int(buf.shape[1]) != self.num_codebooks: + buf = torch.arange(self.num_codebooks, dtype=torch.long, device=device).view(1, self.num_codebooks) + self._codebook_index_cache[key] = buf + return buf + + def _get_boc_frames(self, num_rows: int, device: torch.device) -> torch.Tensor: + key = self._device_cache_key(device) + buf = self._boc_frame_cache.get(key) + if buf is None or buf.device != device or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + buf = torch.full((rows, self.num_codebooks), BOC_ID, dtype=torch.long, device=device) + self._boc_frame_cache[key] = buf + return buf[:num_rows] + + # ------------------------------------------------------------------ postprocess + def postprocess( + self, + hidden_states_slice: torch.Tensor, + multimodal_outputs: Any = None, + **req_infos: Any, + ) -> dict[str, Any]: + """Publish per-request audio codes into model_intermediate_buffer. + + Called once per request in batch order. Indexes _last_audio_codes + by a running cursor (one row per request per step). + """ + _ = multimodal_outputs + return self._postprocess_impl() + + def _postprocess_impl(self) -> dict[str, Any]: + host_staging = self._last_audio_host_staging + codes_full = self._last_audio_codes + if host_staging is None and codes_full is None: + return {} + event = self._last_audio_staging_event + if event is not None: + event.synchronize() + self._last_audio_staging_event = None + + cursor = int(self._postprocess_cursor) + num_rows = int(host_staging.shape[0]) if host_staging is not None else int(codes_full.shape[0]) + if cursor >= num_rows: + self._postprocess_cursor = 0 + return {} + self._postprocess_cursor = cursor + 1 + + if host_staging is not None: + if self._last_audio_valid_flags is None or self._last_audio_done_flags is None: + self._last_audio_valid_flags = host_staging[:num_rows, self.num_codebooks].tolist() + self._last_audio_done_flags = host_staging[:num_rows, self.num_codebooks + 1].tolist() + valid_flag = int(self._last_audio_valid_flags[cursor]) + done_flag = int(self._last_audio_done_flags[cursor]) + if valid_flag or done_flag: + self._postprocess_audio_active_rows += 1 + if cursor + 1 >= num_rows: + if self._postprocess_audio_rows == num_rows and self._postprocess_audio_active_rows == num_rows: + self._fast_audio_direct_rows = num_rows + self._postprocess_audio_rows = 0 + self._postprocess_audio_active_rows = 0 + if valid_flag == 0: + return {} + new_codes = host_staging[cursor : cursor + 1, : self.num_codebooks] + return {"codes": {"audio": new_codes}} + + if cursor >= len(self._last_audio_code_valid) or not self._last_audio_code_valid[cursor]: + return {} + slice_codes = codes_full[cursor : cursor + 1] + new_codes = slice_codes.to(torch.int32) + return {"codes": {"audio": new_codes}} + + # ------------------------------------------------------------------ helpers + def _audio_codebook_logits_from_rows( + self, hidden_states: torch.Tensor, audio_rows: torch.Tensor, *, all_rows: bool = False + ) -> torch.Tensor: + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + if audio_rows.numel() == 0: + return torch.empty( + (0, self.num_codebooks, self.codebook_size), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + if all_rows: + return self.modality_head.generate(hidden_flat[: int(audio_rows.numel())]) + return self.modality_head.generate(hidden_flat.index_select(0, audio_rows.to(hidden_flat.device))) + + def _apply_delay_pattern_masking_batched( + self, cb_logits: torch.Tensor, audio_rows: torch.Tensor, *, all_rows: bool = False + ) -> None: + """Vectorized delay-pattern masking using GPU-resident sampler state.""" + if cb_logits.numel() == 0: + return + + rows = audio_rows.to(device=cb_logits.device, dtype=torch.long) + num_audio_rows = int(cb_logits.shape[0]) + num_codebooks = self.num_codebooks + q = self._get_codebook_indices(cb_logits.device) + if all_rows: + delay = self._decode_delay_count[:num_audio_rows] + rem = self._decode_eoc_countdown[:num_audio_rows] + else: + delay = self._decode_delay_count.index_select(0, rows) + rem = self._decode_eoc_countdown.index_select(0, rows) + if delay.dtype != torch.long: + delay = delay.to(torch.long) + if rem.dtype != torch.long: + rem = rem.to(torch.long) + ramp = rem >= 0 + + bos = BOC_ID + eos = EOC_ID + + # Ramp-down: codebooks already behind the EOC wave are locked to EOC. + lock_until = num_codebooks - rem + locked = ramp.unsqueeze(1) & (q < lock_until.unsqueeze(1)) + eoc_logits = cb_logits[:, :, eos].clone() + cb_logits.masked_fill_(locked.unsqueeze(-1), float("-inf")) + cb_logits[:, :, eos] = torch.where(locked, eoc_logits, cb_logits[:, :, eos]) + + # Ramp-down tail: active codebooks cannot emit stream specials. + ramp_tail = ramp.unsqueeze(1) & (q >= lock_until.unsqueeze(1)) + cb_logits[:, :, bos] = torch.where( + ramp_tail, + torch.full_like(cb_logits[:, :, bos], float("-inf")), + cb_logits[:, :, bos], + ) + cb_logits[:, :, eos] = torch.where( + ramp_tail, + torch.full_like(cb_logits[:, :, eos], float("-inf")), + cb_logits[:, :, eos], + ) + + # Normal generation: delayed codebooks are forced to BOC. + normal = ~ramp + allowed_until = torch.clamp(delay + 1, max=num_codebooks) + delayed = normal.unsqueeze(1) & (q >= allowed_until.unsqueeze(1)) + boc_logits = cb_logits[:, :, bos].clone() + cb_logits.masked_fill_(delayed.unsqueeze(-1), float("-inf")) + cb_logits[:, :, bos] = torch.where(delayed, boc_logits, cb_logits[:, :, bos]) + + # Normal active codebooks: BOC is disallowed; EOC is only allowed on cb0. + allowed = normal.unsqueeze(1) & (q < allowed_until.unsqueeze(1)) + cb_logits[:, :, bos] = torch.where( + allowed, + torch.full_like(cb_logits[:, :, bos], float("-inf")), + cb_logits[:, :, bos], + ) + nonzero_allowed = allowed & (q > 0) + cb_logits[:, :, eos] = torch.where( + nonzero_allowed, + torch.full_like(cb_logits[:, :, eos], float("-inf")), + cb_logits[:, :, eos], + ) + + def _update_delay_state_batched( + self, + codes_flat: torch.Tensor, + audio_rows: torch.Tensor, + num_rows: int, + device: torch.device, + *, + code_row_mask: torch.Tensor | None = None, + all_rows: bool = False, + ) -> None: + """Update delay/ramp-down state in batch and stage per-request outputs.""" + self._ensure_decode_state_capacity(num_rows, device) + if codes_flat.numel() == 0: + self._clear_last_audio_outputs() + return + + rows = audio_rows.to(device=device, dtype=torch.long) + num_audio_rows = int(codes_flat.shape[0]) + num_codebooks = self.num_codebooks + q = self._get_codebook_indices(device) + + if all_rows: + prev_delay = self._decode_delay_count[:num_audio_rows] + prev_rem = self._decode_eoc_countdown[:num_audio_rows] + prev_done = self._decode_generation_done[:num_audio_rows].to(torch.bool) + prev_has_codes = self._decode_has_codes[:num_audio_rows].to(torch.bool) + prev_codes = self._decode_last_codes[:num_audio_rows] + else: + prev_delay = self._decode_delay_count.index_select(0, rows) + prev_rem = self._decode_eoc_countdown.index_select(0, rows) + prev_done = self._decode_generation_done.index_select(0, rows).to(torch.bool) + prev_has_codes = self._decode_has_codes.index_select(0, rows).to(torch.bool) + prev_codes = self._decode_last_codes.index_select(0, rows) + if prev_delay.dtype != torch.long: + prev_delay = prev_delay.to(torch.long) + if prev_rem.dtype != torch.long: + prev_rem = prev_rem.to(torch.long) + if code_row_mask is None: + update_mask = torch.ones_like(prev_has_codes) + elif all_rows: + update_mask = code_row_mask.to(device=device, dtype=torch.bool)[:num_audio_rows] + else: + update_mask = code_row_mask.to(device=device, dtype=torch.bool).index_select(0, rows) + ramp = prev_rem >= 0 + + codes = codes_flat.to(device=device, dtype=torch.long).clone() + + # Leading BOC delay pad. + delay_active = (~ramp) & ((prev_delay + 1) < num_codebooks) + delay_mask = delay_active.unsqueeze(1) & (q > prev_delay.unsqueeze(1)) + codes = torch.where(delay_mask, torch.full_like(codes, BOC_ID), codes) + next_delay = torch.where(delay_active, prev_delay + 1, prev_delay) + + # Trailing EOC ramp-down. + lock_until = num_codebooks - prev_rem + ramp_mask = ramp.unsqueeze(1) & (q < lock_until.unsqueeze(1)) + codes = torch.where(ramp_mask, torch.full_like(codes, EOC_ID), codes) + ramp_next_rem = prev_rem - 1 + + # New EOC detection in normal mode. This avoids tensor->Python control + # flow; ties are impossible for the index because q is monotonic. + eos_mask = (~ramp).unsqueeze(1) & (codes == EOC_ID) + eos_idx_values = torch.where(eos_mask, q.expand_as(codes), torch.full_like(codes, -1)) + last_eos_idx = eos_idx_values.max(dim=1).values + has_eos = last_eos_idx >= 0 + eos_prefix = has_eos.unsqueeze(1) & (q <= last_eos_idx.unsqueeze(1)) + codes = torch.where(eos_prefix, torch.full_like(codes, EOC_ID), codes) + normal_next_rem = torch.where(has_eos, num_codebooks - last_eos_idx - 1, torch.full_like(last_eos_idx, -1)) + + next_rem = torch.where(ramp, ramp_next_rem, normal_next_rem) + done = ((next_rem >= 0) & (next_rem <= 0)) & update_mask + valid = update_mask + + next_delay = torch.where(done, torch.zeros_like(next_delay), next_delay) + next_rem = torch.where(done, torch.full_like(next_rem, -1), next_rem) + output_codes = codes + + write_delay = torch.where(update_mask, next_delay, prev_delay) + write_rem = torch.where(update_mask, next_rem, prev_rem) + write_done = torch.where(update_mask, done, prev_done) + write_codes = torch.where(update_mask.unsqueeze(1), codes, prev_codes) + write_has_codes = torch.where(update_mask, ~done, prev_has_codes) + + if all_rows: + self._decode_delay_count[:num_audio_rows].copy_(write_delay.to(dtype=self._decode_delay_count.dtype)) + self._decode_eoc_countdown[:num_audio_rows].copy_(write_rem.to(dtype=self._decode_eoc_countdown.dtype)) + self._decode_generation_done[:num_audio_rows].copy_(write_done) + self._decode_last_codes[:num_audio_rows].copy_(write_codes) + self._decode_has_codes[:num_audio_rows].copy_(write_has_codes) + else: + self._decode_delay_count.index_copy_(0, rows, write_delay.to(dtype=self._decode_delay_count.dtype)) + self._decode_eoc_countdown.index_copy_(0, rows, write_rem.to(dtype=self._decode_eoc_countdown.dtype)) + self._decode_generation_done.index_copy_(0, rows, write_done) + self._decode_last_codes.index_copy_(0, rows, write_codes) + self._decode_has_codes.index_copy_(0, rows, write_has_codes) + + staged_codes = torch.where(update_mask.unsqueeze(1), output_codes, torch.full_like(output_codes, -1)) + if all_rows: + valid_full = valid + done_full = done + staging = self._get_audio_gpu_staging_buffer(num_rows, device) + staging[:, :num_codebooks].copy_(staged_codes.to(torch.int32)) + staging[:, num_codebooks].copy_(valid_full.to(torch.int32)) + staging[:, num_codebooks + 1].copy_(done_full.to(torch.int32)) + else: + codes_full = self._get_audio_codes_buffer(num_rows, device) + codes_full.fill_(-1) + codes_full.index_copy_(0, rows, staged_codes) + valid_full = torch.zeros(num_rows, dtype=torch.bool, device=device) + done_full = torch.zeros(num_rows, dtype=torch.bool, device=device) + valid_full.index_copy_(0, rows, valid) + done_full.index_copy_(0, rows, done) + staging = self._get_audio_gpu_staging_buffer(num_rows, device) + staging[:, :num_codebooks].copy_(codes_full.to(torch.int32)) + staging[:, num_codebooks].copy_(valid_full.to(torch.int32)) + staging[:, num_codebooks + 1].copy_(done_full.to(torch.int32)) + + host = self._get_audio_host_staging_buffer(num_rows) + host.copy_(staging, non_blocking=True) + if device.type == "cuda": + event = self._next_audio_staging_event() + event.record(torch.cuda.current_stream(device)) + self._last_audio_staging_event = event + else: + self._last_audio_staging_event = None + + self._decode_active_audio_count = num_rows if all_rows else int(rows.numel()) + # Host staging is the authoritative postprocess path. Keeping another + # full GPU codes buffer alive adds fill/copy work on the hot path. + self._last_audio_codes = None + self._last_audio_host_staging = host[:num_rows] + self._last_audio_valid_flags = None + self._last_audio_done_flags = None + self._last_audio_code_valid = [] + self._postprocess_cursor = 0 + self._postprocess_audio_rows = num_rows + self._postprocess_audio_active_rows = 0 + + def _sample_audio_codes(self, logits_2d: torch.Tensor) -> torch.Tensor: + """Replicate upstream sampling: temperature → top-k → top-p → multinomial.""" + x = logits_2d.float() + top_k = 50 + top_p = 0.95 + if 0 < top_k < x.shape[-1]: + kth = x.topk(top_k, dim=-1).values[..., -1:] + x = x.masked_fill(x < kth, float("-inf")) + if 0.0 < top_p < 1.0: + sorted_logits, sorted_idx = x.sort(dim=-1, descending=True) + cumprobs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) + sorted_mask = cumprobs > top_p + sorted_mask[..., 1:] = sorted_mask[..., :-1].clone() + sorted_mask[..., 0] = False + mask = torch.zeros_like(x, dtype=torch.bool) + mask.scatter_(-1, sorted_idx, sorted_mask) + x = x.masked_fill(mask, float("-inf")) + # Detect all-masked rows BEFORE softmax (softmax of all-inf yields NaN). + has_finite = torch.isfinite(x).any(dim=-1) + all_masked = ~has_finite + # Do not branch on tensor bools here; that synchronizes every audio step. + fallback = x.argmax(dim=-1) + safe_x = torch.where(all_masked.unsqueeze(-1), torch.zeros_like(x), x) + probs = safe_x.softmax(dim=-1) + sampled = torch.multinomial(probs, num_samples=1).squeeze(-1) + return torch.where(all_masked, fallback, sampled) + + # ------------------------------------------------------------------ omni output + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + hidden = model_outputs + + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") + if info_dicts is None: + info_dicts = [] + + audio_codes_list: list[torch.Tensor] = [] + any_nonempty = False + for info in info_dicts: + ac: torch.Tensor | None = None + if isinstance(info, dict): + codes_field = info.get("codes") + if isinstance(codes_field, dict): + ac = codes_field.get("audio") + else: + ac = info.get("audio_codes") + if isinstance(ac, torch.Tensor) and ac.numel() > 0: + audio_codes_list.append(ac) + any_nonempty = True + else: + audio_codes_list.append(torch.empty(0, dtype=torch.long)) + + if any_nonempty: + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={"codes": {"audio": audio_codes_list}}, + ) + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=None) + + # ------------------------------------------------------------------ weight loading + + # Per-layer suffixes from the actual V3 checkpoint (results/higgs_v3_checkpoint_analysis.txt) + _V3_LAYER_SUFFIXES = ( + "input_layernorm.weight", + "mlp.down_proj.weight", + "mlp.gate_proj.weight", + "mlp.up_proj.weight", + "post_attention_layernorm.weight", + "self_attn.k_norm.weight", + "self_attn.k_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_norm.weight", + "self_attn.q_proj.weight", + "self_attn.v_proj.weight", + ) + + @classmethod + def _build_required_keys(cls, num_layers: int) -> set[str]: + """Build the exact set of required V3 checkpoint keys.""" + keys = { + "tied.embedding.text_embedding.weight", + "body.norm.weight", + f"{_MODALITY_EMBEDDING_PREFIX}weight", + } + for i in range(num_layers): + for suffix in cls._V3_LAYER_SUFFIXES: + keys.add(f"body.layers.{i}.{suffix}") + return keys + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + backbone_weights: list[tuple[str, torch.Tensor]] = [] + loaded_params: set[str] = set() + own_params = dict(self.named_parameters()) + seen_checkpoint_keys: set[str] = set() + + for name, tensor in weights: + seen_checkpoint_keys.add(name) + + mapped = self._map_weight_name(name) + if mapped is None: + continue + + if mapped.startswith("model.") or mapped.startswith("lm_head."): + backbone_weights.append((mapped, tensor)) + elif mapped in own_params: + param = own_params[mapped] + if param.shape != tensor.shape: + raise ValueError( + f"Shape mismatch for {mapped}: expected {param.shape}, " + f"got {tensor.shape} (checkpoint key: {name})" + ) + param.data.copy_(tensor.to(param.dtype)) + loaded_params.add(mapped) + + if backbone_weights: + backbone_module = _BackboneWrapper(self.model, self.lm_head, self._backbone_config) + loaded = backbone_module.load_weights(iter(backbone_weights)) + loaded_params.update(loaded) + + # Resolve special token IDs from the tokenizer + model_path = getattr(self.vllm_config.model_config, "model", None) + if model_path: + self.config.resolve_special_tokens(model_path) + self._resolve_token_ids() + + # Verify every required checkpoint key was seen. + num_layers = int(self._backbone_config.num_hidden_layers) + required = self._build_required_keys(num_layers) + missing = required - seen_checkpoint_keys + if missing: + raise RuntimeError( + f"HiggsAudioV3Talker: {len(missing)} required checkpoint keys missing: {sorted(missing)[:5]}..." + ) + + logger.info( + "HiggsAudioV3Talker: loaded %d params, modality_embedding=%s, tied=%s", + len(loaded_params), + tuple(self.multimodal_embedding.weight.shape), + self.tie_modality, + ) + return loaded_params + + def _map_weight_name(self, name: str) -> str | None: + if name.startswith(_CODEC_PREFIX): + return None + if name.startswith(_MODALITY_HEAD_PREFIX): + if self.tie_modality: + return None + return name.replace(_MODALITY_HEAD_PREFIX, "modality_head.") + if name.startswith(_MODALITY_EMBEDDING_PREFIX): + return name.replace(_MODALITY_EMBEDDING_PREFIX, "multimodal_embedding.") + for ckpt_prefix, model_prefix in _BACKBONE_PREFIX_MAP.items(): + if name.startswith(ckpt_prefix): + return name.replace(ckpt_prefix, model_prefix, 1) + # Reject unexpected non-codec Higgs checkpoint keys + raise ValueError( + f"Unexpected checkpoint key with no known mapping: {name!r}. " + f"Known prefixes: {list(_BACKBONE_PREFIX_MAP.keys())}, " + f"{_MODALITY_EMBEDDING_PREFIX!r}, {_MODALITY_HEAD_PREFIX!r}, {_CODEC_PREFIX!r}" + ) + + +class _BackboneWrapper(nn.Module): + """Wrapper to use AutoWeightsLoader for Qwen3 backbone.""" + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__(self, model, lm_head, config): + super().__init__() + self.model = model + self.lm_head = lm_head + self.config = config + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + from vllm.model_executor.models.utils import AutoWeightsLoader + + skip = ["lm_head."] if getattr(self.config, "tie_word_embeddings", False) else None + loader = AutoWeightsLoader(self, skip_prefixes=skip) + return loader.load_weights(weights) diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_tokenizer.py new file mode 100644 index 000000000..8f881564a --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_tokenizer.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prompt builder for higgs-audio v3 TTS. + +Prompt formats: + Zero-shot: + <|tts|> <|text|> {text tokens} <|audio|> + Voice clone (no ref text): + <|tts|> <|ref_audio|> [-100]×N <|text|> {text tokens} <|audio|> + Voice clone (with ref text): + <|tts|> <|ref_text|> {ref text tokens} <|ref_audio|> [-100]×N <|text|> {text tokens} <|audio|> + +-100 placeholders are replaced at prefill time with fused multi-codebook +embeddings of the delay-pattern-encoded reference audio codes. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch + +__all__ = [ + "HiggsAudioV3TokenizerAdapter", + "AUDIO_PLACEHOLDER_ID", + "encode_reference_audio", + "apply_delay_pattern", +] + +AUDIO_PLACEHOLDER_ID = -100 + +_REQUIRED_SPECIALS: tuple[str, ...] = ( + "<|tts|>", + "<|text|>", + "<|audio|>", +) + +# Codec stream special token IDs (inside the codebook vocab, not the LM vocab) +BOC_ID = 1024 +EOC_ID = 1025 + + +def apply_delay_pattern(codes_tn: torch.Tensor) -> torch.Tensor: + """Apply MusicGen-style delay pattern to raw codes. + + Input: [T, N] raw codes (T frames, N codebooks). + Output: [T + N - 1, N] delayed codes with BOC/EOC padding. + + Codebook c is delayed by c positions: rows 0..c-1 get BOC, + rows c..c+T-1 get real codes, rows c+T..T+N-2 get EOC. + """ + T, N = codes_tn.shape + L = T + N - 1 + delayed = torch.full((L, N), EOC_ID, dtype=codes_tn.dtype) + for c in range(N): + delayed[:c, c] = BOC_ID + delayed[c : c + T, c] = codes_tn[:, c] + return delayed + + +def encode_reference_audio( + wav: np.ndarray, + sr: int, +) -> torch.Tensor: + """Encode a reference audio clip to codec codes [T, num_codebooks]. + + Uses the same HiggsAudioV2TokenizerModel as v2 (same codec). + Returns raw codes before delay pattern application. + """ + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( + _encode_ref_audio_codes, + ) + + # _encode_ref_audio_codes returns [num_codebooks, T_raw] + codes_qt = _encode_ref_audio_codes(wav, sr) + # Transpose to [T_raw, num_codebooks] + return codes_qt.transpose(0, 1).contiguous() + + +class HiggsAudioV3TokenizerAdapter: + """Wraps the HF tokenizer and builds TTS prompts.""" + + def __init__(self, tokenizer: Any) -> None: + self._tok = tokenizer + vocab = dict(tokenizer.get_added_vocab()) + missing = [t for t in _REQUIRED_SPECIALS if t not in vocab] + if missing: + raise ValueError(f"Tokenizer is missing Higgs TTS v3 specials: {missing}") + self.tts_id: int = vocab["<|tts|>"] + self.text_id: int = vocab["<|text|>"] + self.audio_id: int = vocab["<|audio|>"] + self.ref_audio_id: int | None = vocab.get("<|ref_audio|>") + self.ref_text_id: int | None = vocab.get("<|ref_text|>") + + @property + def tokenizer(self) -> Any: + return self._tok + + def build_prompt( + self, + text: str, + *, + num_ref_tokens: int = 0, + reference_text: str | None = None, + ) -> list[int]: + """Build a TTS prompt. + + Args: + text: Target text to synthesize. + num_ref_tokens: Number of delay-pattern reference code rows. + 0 means zero-shot (no voice clone). + reference_text: Optional transcript of the reference audio. + """ + if not text or not text.strip(): + raise ValueError("Text input must be non-empty for TTS") + if num_ref_tokens < 0: + raise ValueError(f"num_ref_tokens must be >= 0, got {num_ref_tokens}") + + ids: list[int] = [self.tts_id] + + # Voice clone: optional ref text + ref audio placeholders + if reference_text and num_ref_tokens > 0 and self.ref_text_id is not None: + ids.append(self.ref_text_id) + ids.extend(self._tok.encode(reference_text, add_special_tokens=False)) + if num_ref_tokens > 0: + if self.ref_audio_id is None: + raise ValueError("Tokenizer missing <|ref_audio|> for voice clone") + ids.append(self.ref_audio_id) + ids.extend([AUDIO_PLACEHOLDER_ID] * num_ref_tokens) + + ids.append(self.text_id) + ids.extend(self._tok.encode(text, add_special_tokens=False)) + ids.append(self.audio_id) + return ids diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py b/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py new file mode 100644 index 000000000..2f35c497e --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""higgs-audio v3 pipeline: Talker (text -> 8-codebook codec) -> Code2Wav (codec -> 24 kHz PCM). + +Two delivery modes are wired here: + +* ``sync_process_input_func`` runs when the deploy YAML has + ``async_chunk: false``. The orchestrator collects the entire Stage-0 emit, + reverts the delay pattern once, and hands a single payload to Stage 1. +* ``async_chunk_process_next_stage_input_func`` runs when the deploy YAML + has ``async_chunk: true``. Stage 0 dispatches per AR step; the streaming + adapter buffers raw delay-pattern rows, slides a window with left context + and right holdback, and emits codec-ready frames per chunk. Stage 1 + trims the overlap on both ends so the client sees a coherent PCM stream. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.higgs_audio_v3" + +HIGGS_AUDIO_V3_PIPELINE = PipelineConfig( + model_type="higgs_multimodal_qwen3", + model_arch="HiggsAudioV3TalkerForConditionalGeneration", + hf_architectures=("HiggsMultimodalQwen3ForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="higgs_audio_v3", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + # stop_token_ids: the model-owned sampler forces eos at ramp-down + # completion. Safety stops from the actual V3 checkpoint: + # 151643 = <|endoftext|> (eos_token_id from config.json) + # 151671 = <|audio_end|> (audio generation end marker) + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [151643, 151671], + }, + async_chunk_process_next_stage_input_func=f"{_PROC}.talker2code2wav_async_chunk", + ), + StagePipelineConfig( + stage_id=1, + model_stage="code2wav", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="HiggsAudioV3Code2WavForConditionalGeneration", + sync_process_input_func=f"{_PROC}.talker2code2wav", + sampling_constraints={"detokenize": True}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py b/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py new file mode 100644 index 000000000..2a21aafcd --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bounded LRU cache for Higgs Audio v3 reference-audio codec encoding. + +Voice-clone requests that re-send the same ref_audio are otherwise re-encoded +through the DAC encoder on every call. The cache keys by a stable content +hash (size + sample rate + head/mid/tail sentinel bytes) so the same audio +served via different URLs or repeated base64 payloads also hits warm. +Tensors are cloned on get/put so callers and the cache cannot alias storage. +""" + +from __future__ import annotations + +import hashlib +import threading +from collections import OrderedDict +from collections.abc import Callable + +import numpy as np +import torch + +__all__ = ["cached_encode_reference_audio", "cache_stats", "clear_cache"] + +_DEFAULT_MAX_ENTRIES = 128 +_SENTINEL_BYTES = 4096 + +_cache: OrderedDict[str, torch.Tensor] = OrderedDict() +_cache_lock = threading.Lock() +_hits = 0 +_misses = 0 + + +def _content_hash(wav: np.ndarray, sr: int) -> str: + """Stable content hash. Head/mid/tail sentinel keeps cost O(1).""" + raw = wav.tobytes() + n = len(raw) + h = hashlib.sha256() + h.update(n.to_bytes(8, "little")) + h.update(int(sr).to_bytes(4, "little")) + if n == 0: + return h.hexdigest() + if n <= _SENTINEL_BYTES * 3: + h.update(raw) + return h.hexdigest() + h.update(raw[:_SENTINEL_BYTES]) + mid = (n - _SENTINEL_BYTES) // 2 + h.update(raw[mid : mid + _SENTINEL_BYTES]) + h.update(raw[-_SENTINEL_BYTES:]) + return h.hexdigest() + + +def cached_encode_reference_audio( + wav: np.ndarray, + sr: int, + encode_fn: Callable[[np.ndarray, int], torch.Tensor], + max_entries: int = _DEFAULT_MAX_ENTRIES, +) -> torch.Tensor: + """LRU-cached wrapper around ``encode_reference_audio``. + + ``encode_fn`` is the actual codec encoder; passed in so this module + stays decoupled from the tokenizer import (avoids circular imports + during stage-input-processor discovery). + + Returns a CPU tensor clone — never the cache's interior storage. + """ + global _hits, _misses + key = _content_hash(wav, sr) + with _cache_lock: + cached = _cache.get(key) + if cached is not None: + _hits += 1 + _cache.move_to_end(key) + return cached.clone() + + # Miss: encode outside the lock so concurrent misses on different + # refs don't serialize on this single mutex. + codes = encode_fn(wav, sr) + codes_cpu = codes.detach().to("cpu").contiguous() + with _cache_lock: + if key in _cache: + # Concurrent miss winner already stored it; refresh LRU position + # and return that copy to avoid divergent tensors per request. + _cache.move_to_end(key) + _misses += 1 # we still paid the encode cost + return _cache[key].clone() + _cache[key] = codes_cpu + if len(_cache) > max_entries: + _cache.popitem(last=False) + _misses += 1 + return codes_cpu.clone() + + +def cache_stats() -> dict: + """Return (hits, misses, size, hit_rate) — for /metrics-style reporting.""" + with _cache_lock: + total = _hits + _misses + return { + "hits": _hits, + "misses": _misses, + "size": len(_cache), + "hit_rate": (_hits / total) if total > 0 else 0.0, + } + + +def clear_cache() -> None: + """Drop all cached entries. Used by tests.""" + global _hits, _misses + with _cache_lock: + _cache.clear() + _hits = 0 + _misses = 0 diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py old mode 100644 new mode 100755 index 7770bee95..8c983cff1 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -1043,6 +1043,9 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc images = mm_data.get("images", []) logger.debug(f"process image count: {len(images)}") batch_feature = image_processor(prompt, images, **tok_kwargs) + vae_generator_seed = mm_kwargs.get("vae_generator_seed") + if vae_generator_seed is not None and images: + batch_feature["vae_generator_seed"] = torch.full((len(images),), int(vae_generator_seed), dtype=torch.long) return batch_feature def _hf_processor_applies_updates( @@ -1080,6 +1083,8 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc config["base_size"] = MultiModalFieldConfig.batched("image") if "ratio_index" in hf_inputs: config["ratio_index"] = MultiModalFieldConfig.batched("image") + if "vae_generator_seed" in hf_inputs: + config["vae_generator_seed"] = MultiModalFieldConfig.batched("image") return config def _get_prompt_updates( @@ -1452,6 +1457,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo HunyuanImage3Inputs: TypeAlias = HunyuanImage3PixelInputs prefer_model_sampler = True + logitsprocs_need_output_token_ids = True # Siglip2 ViT supports data-parallel encoding (mm_encoder_tp_mode="data"): # weights are replicated and the image batch is sharded across TP ranks. @@ -1732,6 +1738,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo # we reconstruct per-image shapes from vae_token_grid_hw below. kwargs.pop("vae_pixel_size", None) vae_token_grid_hw = kwargs.pop("vae_token_grid_hw", None) + vae_generator_seed = kwargs.pop("vae_generator_seed", None) if vit_pixel_values is None or vae_pixel_values is None: return None @@ -1771,6 +1778,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo "vit_spatial_shapes": vit_spatial_shapes, "vae_pixel_values": vae_image_list, "vae_token_grid_hw": vae_token_grid_hw, + "vae_generator_seed": vae_generator_seed, }, ) @@ -1778,6 +1786,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo self, images: torch.Tensor, cfg_factor: int = 1, + generator: torch.Generator | None = None, ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: """ Encode images through VAE encoder. @@ -1797,11 +1806,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo images = images.to(dtype=self.vae.dtype) vae_encode_result = self.vae.encode(images) - # Match HunyuanImage-3's cond encode path: sample the posterior, but - # use a fixed generator so online requests do not consume the global - # RNG and drift across a long-running server. - _cond_vae_gen = torch.Generator(device=images.device).manual_seed(0) - latents = vae_encode_result.latent_dist.sample(_cond_vae_gen) + latents = vae_encode_result.latent_dist.sample(generator) # Apply shift and scaling factors if present if hasattr(config, "shift_factor") and config.shift_factor: @@ -1933,6 +1938,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo vit_pixel_attention_mask = pixel_values["vit_pixel_attention_mask"] vit_spatial_shapes = pixel_values["vit_spatial_shapes"] vae_pixel_values = pixel_values["vae_pixel_values"] + vae_generator_seed = pixel_values.get("vae_generator_seed") + vae_generator = None + if vae_generator_seed is not None and vae_generator_seed.numel() > 0: + vae_generator = torch.Generator(device=vae_pixel_values[0].device).manual_seed( + int(vae_generator_seed.reshape(-1)[0].item()) + ) # Perform ViT encoding vit_embeddings = self._vit_encode(vit_pixel_values, vit_pixel_attention_mask, vit_spatial_shapes) @@ -1941,7 +1952,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo # `reso_group` bucket so shapes are ragged across the image-batch dim. vae_token_embeddings = [] for vae_image_i in vae_pixel_values: - t_i, latents_i = self._vae_encode(vae_image_i.unsqueeze(0), vae_cfg_factor) + t_i, latents_i = self._vae_encode(vae_image_i.unsqueeze(0), vae_cfg_factor, generator=vae_generator) t_emb = self.time_embed(t_i[0]) vae_tokens, _, _ = self.patch_embed(latents_i, t_emb) vae_token_embeddings.append(vae_tokens) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index 08ed9e854..e7b305b1e 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -26,10 +26,10 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig -from .audio_vae import AudioVAE, AudioVAEConfig from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input from .text_processing import segment_and_normalize diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py index 0a9ce0f5c..8c527b8d9 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py @@ -72,6 +72,7 @@ from vllm_omni.transformers_utils.processors.ming import ( PLACEHOLDER_VIDEO_TOKEN_IN_TEXT, MingFlashOmniProcessor, MingWhisperFeatureExtractor, + raise_missing_video_processor, ) from .audio_encoder import WhisperAudioEncoder @@ -540,20 +541,20 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO if images is not None: image_outputs = hf_processor.image_processor( images=images, - videos=None, return_tensors="pt", ) data.update(image_outputs) videos = mm_data.get("videos", None) if videos is not None: - # TODO: ``videos=`` on image_processor is deprecated since - # transformers v4.57 (removed in v5); migrate to Qwen2VLVideoProcessor. - video_outputs = hf_processor.image_processor( - images=None, - videos=videos, - return_tensors="pt", - ) + video_processor = getattr(hf_processor, "video_processor", None) + if video_processor is not None: + video_outputs = video_processor( + videos=videos, + return_tensors="pt", + ) + else: + raise_missing_video_processor() # Rename keys to distinguish from images if "pixel_values" in video_outputs: video_outputs["pixel_values_videos"] = video_outputs.pop("pixel_values") @@ -696,6 +697,9 @@ class MingFlashOmniThinkerForConditionalGeneration( f"image_gen_scales={list(self.query_tokens_dict.keys())}" ) + def get_language_model(self): + return self.language_model + def _load_image_gen_query_tokens(self, model_path: str) -> None: """Load learnable query tokens used for image generation. diff --git a/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py b/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py index ca7d00f50..5a1af2fc2 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py @@ -714,6 +714,9 @@ class BailingMoeV2Model(nn.Module): def get_input_embeddings(self): return self.word_embeddings + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.word_embeddings(input_ids) + def forward( self, input_ids: torch.Tensor, @@ -794,6 +797,9 @@ class BailingMoeV2ForCausalLM(nn.Module, CustomProcessMixin): self.sampler = Sampler() self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + def forward( self, input_ids: torch.Tensor, @@ -818,7 +824,7 @@ class BailingMoeV2ForCausalLM(nn.Module, CustomProcessMixin): hidden_states: torch.Tensor, sampling_metadata, ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states, sampling_metadata) + logits = self.logits_processor(self.lm_head, hidden_states) return logits def sample( diff --git a/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py index 53c130b90..18deb4600 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py @@ -110,12 +110,18 @@ BASE_CAPTION_TEMPLATE: dict[str, Any] = { } -def create_instruction(user_input: dict[str, Any]) -> str: +def create_instruction(user_input: Any) -> str | None: """Return a JSON caption string for ``audio_sequence[0]``. Only keys already present on the base template are merged in; unknown keys are silently ignored to keep the output schema stable. """ + if user_input is None: + return None + if isinstance(user_input, str): + return user_input + if not isinstance(user_input, dict): + raise ValueError(f"Ming instruction must be str, dict, or None; got {type(user_input).__name__}") caption = copy.deepcopy(BASE_CAPTION_TEMPLATE) item = caption["audio_sequence"][0] for key, value in user_input.items(): diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index 6404b3fb8..cfb1090c2 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -27,223 +27,19 @@ from threading import Lock import torch import torch.nn as nn -import torch.nn.functional as F from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, StaticCache from vllm.logger import init_logger from vllm.platforms import current_platform -from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb +from x_transformers.x_transformers import RotaryEmbedding from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding - -from .audio_vae import AudioVAE +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE +from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps +from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps logger = init_logger(__name__) -######################################################################## -# DiT Modules -# Ported from: -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/modules.py -# Ported from: -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/dit.py -######################################################################## - - -class RMSNorm(nn.Module): - def __init__(self, dim: int, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.ones(dim)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.weight.dtype in [torch.float16, torch.bfloat16]: - x = x.to(self.weight.dtype) - x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) - return x - - -class FeedForward(nn.Module): - def __init__( - self, dim: int, dim_out: int | None = None, mult: float = 4, dropout: float = 0.0, approximate: str = "none" - ): - super().__init__() - inner_dim = int(dim * mult) - dim_out = dim_out if dim_out is not None else dim - - activation = nn.GELU(approximate=approximate) - project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) - self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.ff(x) - - -class Attention(nn.Module): - def __init__( - self, - dim: int, - heads: int = 8, - dim_head: int = 64, - dropout: float = 0.0, - qk_norm: str | None = None, - pe_attn_head: int | None = None, - attn_mask_enabled: bool = True, - ): - super().__init__() - self.dim = dim - self.heads = heads - self.inner_dim = dim_head * heads - self.dropout = dropout - - self.to_q = nn.Linear(dim, self.inner_dim) - self.to_k = nn.Linear(dim, self.inner_dim) - self.to_v = nn.Linear(dim, self.inner_dim) - if qk_norm is None: - self.q_norm = None - self.k_norm = None - elif qk_norm == "rms_norm": - self.q_norm = RMSNorm(dim_head) - self.k_norm = RMSNorm(dim_head) - else: - raise ValueError(f"Unimplemented qk_norm: {qk_norm}") - - self.to_out = nn.ModuleList([]) - self.to_out.append(nn.Linear(self.inner_dim, dim)) - self.to_out.append(nn.Dropout(dropout)) - - self.pe_attn_head = pe_attn_head - self.attn_mask_enabled = attn_mask_enabled - - def forward( - self, - x: torch.Tensor, - mask: torch.Tensor | None = None, - rope: tuple[torch.Tensor, torch.Tensor | None] | None = None, - ) -> torch.Tensor: - batch_size = x.shape[0] - - query = self.to_q(x) - key = self.to_k(x) - value = self.to_v(x) - - inner_dim = key.shape[-1] - head_dim = inner_dim // self.heads - query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) - - if self.q_norm is not None: - query = self.q_norm(query) - if self.k_norm is not None: - key = self.k_norm(key) - - if rope is not None: - freqs, xpos_scale = rope - q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) - - if self.pe_attn_head is not None: - on = self.pe_attn_head - query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale) - key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale) - else: - query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) - key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) - - if self.attn_mask_enabled and mask is not None: - valid_sample_indices = mask.any(dim=1) - final_output = torch.zeros_like(query).to(query.device) - - attn_mask = mask[valid_sample_indices] - query = query[valid_sample_indices] - key = key[valid_sample_indices] - value = value[valid_sample_indices] - attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) - attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2]) - else: - attn_mask = None - - x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) - if self.attn_mask_enabled and mask is not None: - final_output[valid_sample_indices] = x - x = final_output - - x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) - x = x.to(query.dtype) - - x = self.to_out[0](x) - x = self.to_out[1](x) - - if mask is not None: - mask = mask.unsqueeze(-1) - x = x.masked_fill(~mask, 0.0) - - return x - - -class DiTBlock(nn.Module): - """A DiT block with pre-norm and residual connections.""" - - def __init__( - self, - hidden_size: int, - num_heads: int, - mlp_ratio: float = 4.0, - dropout: float = 0.1, - qk_norm: str | None = None, - pe_attn_head: int | None = None, - attn_mask_enabled: bool = True, - **kwargs, - ): - super().__init__() - self.norm1 = RMSNorm(hidden_size) - self.attn = Attention( - dim=hidden_size, - heads=num_heads, - dim_head=hidden_size // num_heads, - dropout=dropout, - qk_norm=qk_norm, - pe_attn_head=pe_attn_head, - attn_mask_enabled=attn_mask_enabled, - ) - self.norm2 = RMSNorm(hidden_size) - self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh") - - def forward( - self, - x: torch.Tensor, - mask: torch.Tensor | None, - rope: tuple[torch.Tensor, torch.Tensor | None] | None, - ) -> torch.Tensor: - x = x + self.attn(self.norm1(x), mask=mask, rope=rope) - x = x + self.mlp(self.norm2(x)) - return x - - -class FinalLayer(nn.Module): - """The final layer of DiT.""" - - def __init__(self, hidden_size: int, out_channels: int): - super().__init__() - self.norm_final = RMSNorm(hidden_size) - self.linear = nn.Linear(hidden_size, out_channels, bias=True) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.norm_final(x) - x = self.linear(x) - return x - - -class CondEmbedder(nn.Module): - """Embeds LLM hidden states with optional CFG dropout.""" - - def __init__(self, input_feature_size: int, hidden_size: int): - super().__init__() - self.cond_embedder = nn.Linear(input_feature_size, hidden_size) - - def forward(self, llm_cond: torch.Tensor) -> torch.Tensor: - return self.cond_embedder(llm_cond) - - class DiT(nn.Module): """Diffusion model with a Transformer backbone for audio latent generation.""" @@ -324,29 +120,6 @@ class DiT(nn.Module): return model_out[:, -x.shape[1] :, :] -######################################################################################### -# CFM -# Ported from: -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/cfm.py -######################################################################################### - - -def get_epss_timesteps(n, device, dtype): - dt = 1 / 32 - predefined_timesteps = { - 5: [0, 2, 4, 8, 16, 32], - 6: [0, 2, 4, 6, 8, 16, 32], - 7: [0, 2, 4, 6, 8, 16, 24, 32], - 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], - 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], - 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], - } - t = predefined_timesteps.get(n, []) - if not t: - return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) - return dt * torch.tensor(t, device=device, dtype=dtype) - - class CFM(nn.Module): """Conditional Flow Matching module for audio latent generation.""" @@ -391,15 +164,8 @@ class CFM(nn.Module): pred, null_pred = torch.chunk(pred_cfg, 2, dim=0) return pred + (pred - null_pred) * sde_args[0] - if self.sway_sampling_coef is not None: - t = t + self.sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) - - for step in range(self.steps): - dt = t[step + 1] - t[step] - y0 = y0 + fn(t[step], y0) * dt - y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step] - - return y0 + t = apply_sway_sampling(t, self.sway_sampling_coef) + return integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, self.steps) class CFMGraphExecutor: @@ -902,6 +668,7 @@ class MingAudioGenerator: if self._audio_vae is None: return requested_max_steps + # Transformers >=5.x may expose these config values as 0-d tensors. sample_rate = float(self._audio_vae.config.sample_rate) vae_patch_size = float(getattr(self._audio_vae.config, "patch_size", 4)) hop_size = float(getattr(self._audio_vae.decoder, "hop_length", 320)) @@ -1041,7 +808,7 @@ class MingAudioGenerator: use_cache=True, ) else: - past_seen_tokens = past_key_values.get_seq_length() + past_seen_tokens = int(past_key_values.get_seq_length()) cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], diff --git a/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py b/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py index 5f54687c0..0a096df86 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py @@ -22,7 +22,8 @@ from .spk_embedding import SpkembExtractor from .talker_module import resample if TYPE_CHECKING: - from .audio_vae import AudioVAE + from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE + from .talker_module import Aggregator logger = init_logger(__name__) diff --git a/vllm_omni/model_executor/models/ming_tts/__init__.py b/vllm_omni/model_executor/models/ming_tts/__init__.py new file mode 100644 index 000000000..35afa2ce5 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .config_ming_tts import MingDenseConfig +from .ming_tts import MingTTSForConditionalGeneration +from .ming_tts_audio_vae import MingAudioVAEModel +from .ming_tts_llm import MingLLMModel + +__all__ = [ + "MingDenseConfig", + "MingTTSForConditionalGeneration", + "MingLLMModel", + "MingAudioVAEModel", +] diff --git a/vllm_omni/model_executor/models/ming_tts/aggregator.py b/vllm_omni/model_executor/models/ming_tts/aggregator.py new file mode 100644 index 000000000..3233c9659 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/aggregator.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/dit.py +from __future__ import annotations + +import torch +import torch.nn as nn +from x_transformers.x_transformers import RotaryEmbedding + +from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer + + +class Aggregator(nn.Module): + def __init__( + self, + in_channels=4, + hidden_size=1152, + depth=28, + num_heads=16, + mlp_ratio=4.0, + llm_input_dim=896, + **kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.num_heads = num_heads + + self.word_embedder = nn.Embedding(1, hidden_size) + self.x_embedder = nn.Linear(in_channels, hidden_size) + self.hidden_size = hidden_size + self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) + self.blocks = nn.ModuleList( + [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] + ) + self.final_layer = FinalLayer(hidden_size, llm_input_dim) + + def forward(self, x, mask=None): + if x.ndim != 3: + raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}") + if x.shape[-1] != self.in_channels: + raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}") + + # [Batch, Time, Dimension] -> [Batch, Time, Hidden]. + x = self.x_embedder(x) + cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device)) + # Prepend a learned CLS token: [Batch, Time, Hidden] -> [Batch, Time + 1, Hidden]. + x = torch.cat([cls_embed, x], dim=1) + + rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) + if mask is not None: + if mask.ndim != 2: + raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}") + if mask.shape[0] != x.shape[0] or mask.shape[1] != x.shape[1] - 1: + raise ValueError( + f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x.shape[0], x.shape[1] - 1)}" + ) + mask_pad = mask.clone().detach()[:, :1] + mask = torch.cat([mask_pad, mask], dim=-1) + for block in self.blocks: + x = block(x, mask, rope) + x = self.final_layer(x) + # Keep the CLS projection only: [Batch, Time + 1, Hidden] -> [Batch, 1, Hidden]. + return x[:, :1, :] + + +__all__ = ["Aggregator"] diff --git a/vllm_omni/model_executor/models/ming_tts/audio_prep.py b/vllm_omni/model_executor/models/ming_tts/audio_prep.py new file mode 100644 index 000000000..b63e70cc8 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/audio_prep.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import math +from io import BytesIO +from typing import Any + +import torch + +from .config_ming_tts import ( + AUDIO_FRAME_HOP, + LATENT_DIM, + PATCH_SIZE, + SAMPLE_RATE, + VAE_PATCH_SIZE, + VISION_START_TOKEN_ID, + MingTTSConfig, +) + + +def pad_prompt_waveform( + waveform: Any, + *, + patch_size: int = PATCH_SIZE, + sample_rate: int = SAMPLE_RATE, +) -> torch.Tensor: + tensor = coerce_prompt_waveform(waveform) + pad_align = int((float(sample_rate) / 12.5) * int(patch_size)) + new_len = ((int(tensor.shape[-1]) + pad_align - 1) // pad_align) * pad_align + if new_len == int(tensor.shape[-1]): + return tensor + padded = torch.zeros((1, new_len), dtype=tensor.dtype, device=tensor.device) + padded[:, : tensor.shape[-1]] = tensor + return padded + + +def coerce_prompt_waveform(value: Any) -> torch.Tensor: + if value is None: + raise ValueError("prompt waveform cannot be None") + if isinstance(value, torch.Tensor): + tensor = value.detach() + if tensor.ndim == 1: + return tensor.unsqueeze(0).to(torch.float32) + if tensor.ndim == 2: + if tensor.shape[0] != 1: + return tensor.reshape(1, -1).to(torch.float32) + return tensor.to(torch.float32) + raise ValueError(f"Unsupported Ming prompt waveform rank: {tuple(tensor.shape)}") + if isinstance(value, (list, tuple)): + parts = [coerce_prompt_waveform(item) for item in value if item is not None] + if not parts: + raise ValueError("prompt waveform list was empty") + return torch.cat(parts, dim=-1) + return coerce_prompt_waveform(torch.as_tensor(value)) + + +def coerce_speaker_embeddings(value: Any, *, use_zero_spk_emb: bool = False) -> list[torch.Tensor] | None: + if value is None: + return [torch.zeros((192,), dtype=torch.float32)] if use_zero_spk_emb else None + if isinstance(value, torch.Tensor): + tensor = value.detach() + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + if tensor.ndim != 2: + raise ValueError(f"Unsupported Ming speaker embedding shape: {tuple(tensor.shape)}") + items = [row.reshape(-1).to(torch.float32).cpu() for row in tensor] + elif isinstance(value, (list, tuple)): + if value and all(not isinstance(item, (list, tuple, torch.Tensor)) for item in value): + items = [torch.as_tensor(value).detach().reshape(-1).to(torch.float32).cpu()] + else: + items = [] + for item in value: + if item is None: + continue + if not isinstance(item, torch.Tensor): + item = torch.as_tensor(item) + items.append(item.detach().reshape(-1).to(torch.float32).cpu()) + else: + return coerce_speaker_embeddings(torch.as_tensor(value), use_zero_spk_emb=use_zero_spk_emb) + if not items: + return [torch.zeros((192,), dtype=torch.float32)] if use_zero_spk_emb else None + for item in items: + if int(item.numel()) != 192: + raise ValueError(f"Ming speaker embedding must have 192 dims, got {int(item.numel())}") + return items + + +def count_prompt_latent_patches( + value: Any, + *, + patch_size: int = PATCH_SIZE, + latent_dim: int = LATENT_DIM, +) -> int: + if value is None: + return 0 + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + latents = value.detach() + if latents.ndim == 3 and latents.shape[0] == 1: + latents = latents.squeeze(0) + if latents.ndim == 3 and latents.shape[-2:] == (patch_size, latent_dim): + return int(latents.shape[0]) + if latents.ndim != 2 or latents.shape[-1] != latent_dim: + raise ValueError(f"Unsupported Ming prompt_latents shape: {tuple(latents.shape)}") + if latents.shape[0] % patch_size != 0: + raise ValueError( + f"Ming prompt_latents frame count must be divisible by patch_size={patch_size}, " + f"got frames={int(latents.shape[0])}" + ) + return int(latents.shape[0] // patch_size) + + +def count_prompt_waveform_patches( + value: Any, + *, + patch_size: int = PATCH_SIZE, + frame_hop: int = AUDIO_FRAME_HOP, + vae_patch_size: int = VAE_PATCH_SIZE, +) -> int: + if value is None: + return 0 + waveform = pad_prompt_waveform(value, patch_size=patch_size) + frame_count = int(math.ceil(float(waveform.shape[-1]) / float(frame_hop))) + latent_frames = int(math.ceil(float(frame_count) / float(vae_patch_size))) + if latent_frames % int(patch_size) != 0: + raise ValueError( + f"Ming prompt waveform produced latent frame count not divisible by patch_size={patch_size}: " + f"frames={latent_frames}" + ) + return int(latent_frames // int(patch_size)) + + +def _normalize_prompt_waveform(value: Any, *, target_sr: int) -> torch.Tensor: + if isinstance(value, bytes): + import torchaudio + + waveform, sr = torchaudio.load(BytesIO(value)) + waveform = waveform[:1].to(torch.float32) + if int(sr) != int(target_sr): + from torchaudio.functional import resample as resample_audio + + waveform = resample_audio(waveform, int(sr), int(target_sr)) + return waveform + + if isinstance(value, tuple) and len(value) == 2 and isinstance(value[1], int): + waveform = coerce_prompt_waveform(value[0]) + if int(value[1]) != int(target_sr): + from torchaudio.functional import resample as resample_audio + + waveform = resample_audio(waveform, int(value[1]), int(target_sr)) + return waveform + + if isinstance(value, dict): + samples = value.get("samples", value.get("array", value.get("waveform"))) + sr = value.get("sample_rate", value.get("sr", target_sr)) + return _normalize_prompt_waveform((samples, int(sr)), target_sr=target_sr) + + return coerce_prompt_waveform(value) + + +def _coerce_prompt_latents( + value: Any, + *, + patch_size: int, + latent_dim: int, +) -> dict[str, torch.Tensor] | None: + if value is None: + return None + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + + latents = value.detach() + if latents.ndim == 3 and latents.shape[0] == 1: + latents = latents.squeeze(0) + + if latents.ndim == 3 and latents.shape[-2:] == (patch_size, latent_dim): + patches = latents + # [Patch, Time, Dimension] -> [Frame, Dimension] for history seeding. + frames = patches.reshape(-1, latent_dim) + return {"patches": patches, "frames": frames} + + if latents.ndim != 2 or latents.shape[-1] != latent_dim: + raise ValueError(f"Unsupported prompt latent shape: {tuple(latents.shape)}") + if latents.shape[0] % patch_size != 0: + raise ValueError( + f"Prompt latent frame count must be divisible by patch_size={patch_size}, " + f"got frames={int(latents.shape[0])}" + ) + # [Frame, Dimension] -> [Patch, Time, Dimension] for Aggregator prompt slots. + patches = latents.reshape(-1, patch_size, latent_dim) if latents.shape[0] > 0 else None + return {"patches": patches, "frames": latents} + + +def _initial_history( + frames: torch.Tensor | None, + *, + history_size: int, + latent_dim: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + history = torch.zeros((history_size, latent_dim), device=device, dtype=dtype) + if frames is None or frames.numel() == 0: + return history + frames = frames.to(device=device, dtype=dtype) + take = min(history_size, int(frames.shape[0])) + history[-take:] = frames[-take:] + return history + + +def _take_scalar(value: Any, idx: int) -> float | None: + if not isinstance(value, torch.Tensor) or value.numel() == 0: + return None + return float(value.reshape(-1)[idx].item()) + + +def _find_audio_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfig) -> torch.Tensor: + dummy_pos = (input_ids == cfg.audio_dummy_token_id).nonzero(as_tuple=True)[0] + if dummy_pos.numel() == 0: + return dummy_pos + + audio_start_pos = (input_ids == cfg.audio_start_token_id).nonzero(as_tuple=True)[0] + audio_end_pos = (input_ids == cfg.audio_end_token_id).nonzero(as_tuple=True)[0] + if audio_start_pos.numel() == 0: + return dummy_pos + + start = int(audio_start_pos[0].item()) + end = int(audio_end_pos[0].item()) if audio_end_pos.numel() > 0 else int(input_ids.shape[0]) + keep = (dummy_pos > start) & (dummy_pos < end) + filtered = dummy_pos[keep] + return filtered if filtered.numel() > 0 else dummy_pos + + +def _find_speaker_placeholder_positions(input_ids: torch.Tensor, hf_config: Any) -> list[int]: + vision_start_token_id = getattr(hf_config, "vision_start_token_id", VISION_START_TOKEN_ID) + vision_start_pos = (input_ids == int(vision_start_token_id)).nonzero(as_tuple=True)[0] + if vision_start_pos.numel() == 0: + return [] + + slots = [] + for pos in vision_start_pos: + slot = int(pos.item()) + 1 + if slot < int(input_ids.shape[0]): + slots.append(slot) + return slots + + +__all__ = [ + "coerce_prompt_waveform", + "coerce_speaker_embeddings", + "count_prompt_latent_patches", + "count_prompt_waveform_patches", + "pad_prompt_waveform", + "_coerce_prompt_latents", + "_find_audio_placeholder_positions", + "_find_speaker_placeholder_positions", + "_initial_history", + "_normalize_prompt_waveform", + "_take_scalar", +] diff --git a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py new file mode 100644 index 000000000..b42c4fdca --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ming dense checkpoint config adapters.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from transformers import PretrainedConfig, Qwen2Config + +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig + +from .constants import ( + AGGREGATOR_HIDDEN_SIZE, + AUDIO_DUMMY_TOKEN_ID, + AUDIO_END_TOKEN_ID, + AUDIO_EOS_TOKEN_ID, + AUDIO_FRAME_HOP, + AUDIO_START_TOKEN_ID, + DEFAULT_CFG, + DEFAULT_SIGMA, + DEFAULT_TEMPERATURE, + HISTORY_PATCH_SIZE, + KEY_CFG, + KEY_CHUNK_ID, + KEY_DECODE_STEP, + KEY_LAST_STOP_PROB, + KEY_LATENT_HISTORY, + KEY_MAX_DECODE_STEPS, + KEY_MIN_DECODE_STEPS, + KEY_NEXT_EMBEDS, + KEY_PROMPT_LATENT_TAIL, + KEY_PROMPT_LATENTS, + KEY_REQUEST_ID, + KEY_SIGMA, + KEY_SPEAKER_EMBEDDING, + KEY_TEMPERATURE, + KEY_TEXT_MODE, + LATENT_CHUNK_SIZE, + LATENT_DIM, + LATENT_LEFT_CONTEXT, + LLM_HIDDEN_SIZE, + LLM_VOCAB_SIZE, + MAX_DECODE_STEPS, + PATCH_SIZE, + SAMPLE_RATE, + STOP_HEAD_MIN_STEPS, + STOP_HEAD_THRESHOLD, + TEXT_EOS_TOKEN_ID, + VAE_PATCH_SIZE, + VISION_START_TOKEN_ID, +) +from .validation import _coerce_audio_vae_config, _nested_get, _to_plain_dict, validate_ming_tts_config + + +def _coerce_qwen2_config(value: Any) -> Qwen2Config: + if isinstance(value, Qwen2Config): + return value + if isinstance(value, PretrainedConfig): + return Qwen2Config.from_dict(value.to_dict()) + if isinstance(value, dict): + return Qwen2Config.from_dict(dict(value)) + raise TypeError(f"Unsupported llm_config type for Ming dense config: {type(value)!r}") + + +class MingDenseConfig(PretrainedConfig): + # The upstream checkpoint declares model_type="dense". Keep it for HF + # config compatibility; deploy/ming_tts.yaml selects the vLLM-Omni + # pipeline via pipeline: ming_tts. + model_type = "dense" + + def __init__( + self, + llm_config: Qwen2Config | dict[str, Any] | None = None, + ditar_config: dict[str, Any] | None = None, + aggregator_config: dict[str, Any] | None = None, + audio_tokenizer_config: AudioVAEConfig | dict[str, Any] | None = None, + architectures: list[str] | None = None, + **kwargs: Any, + ) -> None: + self.llm_config = _coerce_qwen2_config(llm_config or {}) + self.ditar_config = dict(ditar_config or {}) + self.aggregator_config = dict(aggregator_config or {}) + self.audio_tokenizer_config = _coerce_audio_vae_config(audio_tokenizer_config) + super().__init__(architectures=architectures, **kwargs) + + def get_text_config(self, decoder: bool = False, **kwargs: Any) -> Qwen2Config: + del decoder, kwargs + return self.llm_config + + +@dataclass +class MingTTSConfig: + """Flat config object shared by Stage-1 and Stage-2. Build via from_hf_config().""" + + llm_hidden_size: int = LLM_HIDDEN_SIZE + llm_vocab_size: int = LLM_VOCAB_SIZE + llm_config: dict[str, Any] = field(default_factory=dict) + + latent_dim: int = LATENT_DIM + patch_size: int = PATCH_SIZE + history_patch_size: int = HISTORY_PATCH_SIZE + + ditar_config: dict[str, Any] = field(default_factory=dict) + aggregator_config: dict[str, Any] = field(default_factory=dict) + + audio_tokenizer_config: AudioVAEConfig | None = None + vae_patch_size: int = VAE_PATCH_SIZE + sample_rate: int = SAMPLE_RATE + audio_frame_hop: int = AUDIO_FRAME_HOP + + cfg: float = DEFAULT_CFG + sigma: float = DEFAULT_SIGMA + temperature: float = DEFAULT_TEMPERATURE + stop_head_min_steps: int = STOP_HEAD_MIN_STEPS + stop_head_threshold: float = STOP_HEAD_THRESHOLD + max_decode_steps: int = MAX_DECODE_STEPS + + latent_chunk_size: int = LATENT_CHUNK_SIZE + latent_left_context: int = LATENT_LEFT_CONTEXT + + text_eos_token_id: int = TEXT_EOS_TOKEN_ID + audio_dummy_token_id: int = AUDIO_DUMMY_TOKEN_ID + audio_start_token_id: int = AUDIO_START_TOKEN_ID + audio_end_token_id: int = AUDIO_END_TOKEN_ID + audio_eos_token_id: int = AUDIO_EOS_TOKEN_ID + + @classmethod + def from_hf_config(cls, hf_config: PretrainedConfig) -> MingTTSConfig: + llm_raw = getattr(hf_config, "llm_config", {}) or {} + ditar_raw = getattr(hf_config, "ditar_config", {}) or {} + agg_raw = getattr(hf_config, "aggregator_config", {}) or {} + atc_raw = getattr(hf_config, "audio_tokenizer_config", None) + + llm_dict = _to_plain_dict(llm_raw) + ditar = _to_plain_dict(ditar_raw) + agg = _to_plain_dict(agg_raw) + ditar.setdefault("attn_backend", "torch") + + atc = _coerce_audio_vae_config(atc_raw) + atc_enc_latent_dim = _nested_get(atc, "enc_kwargs", "latent_dim", default=LATENT_DIM) + atc_patch_size = _nested_get(atc, "patch_size", default=VAE_PATCH_SIZE) + atc_sample_rate = _nested_get(atc, "sample_rate", default=SAMPLE_RATE) + + enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=AUDIO_FRAME_HOP) + + cfg = cls( + llm_hidden_size=llm_dict.get("hidden_size", LLM_HIDDEN_SIZE), + llm_vocab_size=llm_dict.get("vocab_size", LLM_VOCAB_SIZE), + llm_config=llm_dict, + latent_dim=atc_enc_latent_dim, + patch_size=ditar.get("patch_size", PATCH_SIZE), + history_patch_size=ditar.get("history_patch_size", HISTORY_PATCH_SIZE), + ditar_config=ditar, + aggregator_config=agg, + audio_tokenizer_config=atc, + vae_patch_size=atc_patch_size, + sample_rate=atc_sample_rate, + audio_frame_hop=enc_hop_size if enc_hop_size is not None else AUDIO_FRAME_HOP, + ) + return cfg + + def validate(self) -> None: + validate_ming_tts_config(self) + + +__all__ = [ + "AGGREGATOR_HIDDEN_SIZE", + "AUDIO_DUMMY_TOKEN_ID", + "AUDIO_END_TOKEN_ID", + "AUDIO_EOS_TOKEN_ID", + "AUDIO_FRAME_HOP", + "AUDIO_START_TOKEN_ID", + "DEFAULT_CFG", + "DEFAULT_SIGMA", + "DEFAULT_TEMPERATURE", + "HISTORY_PATCH_SIZE", + "KEY_CFG", + "KEY_CHUNK_ID", + "KEY_DECODE_STEP", + "KEY_LAST_STOP_PROB", + "KEY_LATENT_HISTORY", + "KEY_MAX_DECODE_STEPS", + "KEY_MIN_DECODE_STEPS", + "KEY_NEXT_EMBEDS", + "KEY_PROMPT_LATENT_TAIL", + "KEY_PROMPT_LATENTS", + "KEY_REQUEST_ID", + "KEY_SIGMA", + "KEY_SPEAKER_EMBEDDING", + "KEY_TEMPERATURE", + "KEY_TEXT_MODE", + "LATENT_CHUNK_SIZE", + "LATENT_DIM", + "LATENT_LEFT_CONTEXT", + "LLM_HIDDEN_SIZE", + "LLM_VOCAB_SIZE", + "MAX_DECODE_STEPS", + "MingDenseConfig", + "MingTTSConfig", + "PATCH_SIZE", + "SAMPLE_RATE", + "STOP_HEAD_MIN_STEPS", + "STOP_HEAD_THRESHOLD", + "TEXT_EOS_TOKEN_ID", + "VAE_PATCH_SIZE", + "VISION_START_TOKEN_ID", +] diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py new file mode 100644 index 000000000..2f59e2094 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/constants.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +# --------------------------------------------------------------------------- +# Token IDs (confirmed from tokenizer_config.json) +# --------------------------------------------------------------------------- + +AUDIO_DUMMY_TOKEN_ID = 151705 # +AUDIO_START_TOKEN_ID = 151706 # +AUDIO_EOS_TOKEN_ID = 151704 # +VISION_START_TOKEN_ID = 151652 # <|vision_start|> + +TEXT_EOS_TOKEN_ID = 151669 # + + +# --------------------------------------------------------------------------- +# Architectural constants (confirmed from original config.json) +# --------------------------------------------------------------------------- + +LATENT_DIM = 64 +PATCH_SIZE = 4 +HISTORY_PATCH_SIZE = 32 +LLM_HIDDEN_SIZE = 896 +LLM_VOCAB_SIZE = 151936 +AGGREGATOR_HIDDEN_SIZE = 1024 +VAE_PATCH_SIZE = 4 +SAMPLE_RATE = 44100 + +# AudioVAE frame/hop geometry (confirmed) +AUDIO_FRAME_HOP = 882 # enc input_dim / hop_size / dec output_dim + +# stop_head defaults +STOP_HEAD_MIN_STEPS = 3 +STOP_HEAD_THRESHOLD = 0.5 + +# FlowLoss sampling defaults +DEFAULT_CFG = 2.0 +DEFAULT_SIGMA = 0.25 +DEFAULT_TEMPERATURE = 0.0 + +# Connector / Stage-2 streaming defaults (runtime tuning) +LATENT_CHUNK_SIZE = 25 +LATENT_LEFT_CONTEXT = 0 +MAX_DECODE_STEPS = 200 + + +# --------------------------------------------------------------------------- +# seq_data.extra_data keys +# --------------------------------------------------------------------------- + +KEY_LATENT_HISTORY = "ming_latent_history" +KEY_DECODE_STEP = "ming_decode_step" +KEY_LAST_STOP_PROB = "ming_last_stop_prob" +KEY_NEXT_EMBEDS = "ming_next_embeds" +KEY_PROMPT_LATENTS = "ming_prompt_latents" +KEY_PROMPT_LATENT_TAIL = "ming_prompt_latent_tail" +KEY_SPEAKER_EMBEDDING = "ming_speaker_embedding" +KEY_REQUEST_ID = "ming_request_id" +KEY_CHUNK_ID = "ming_chunk_id" +KEY_CFG = "ming_cfg" +KEY_SIGMA = "ming_sigma" +KEY_TEMPERATURE = "ming_temperature" +KEY_MAX_DECODE_STEPS = "ming_max_decode_steps" +KEY_MIN_DECODE_STEPS = "ming_min_decode_steps" +KEY_TEXT_MODE = "ming_text_mode" diff --git a/vllm_omni/model_executor/models/ming_tts/flowloss_head.py b/vllm_omni/model_executor/models/ming_tts/flowloss_head.py new file mode 100644 index 000000000..3f300cec0 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/flowloss_head.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/flowloss.py + +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.common.ming.fm import Solver, build_timesteps + +from .fm.dit import DiT + + +class FlowLoss(nn.Module): + """Diffusion Loss""" + + def __init__(self, z_channels, llm_cond_dim, **kwargs): + super().__init__() + self.z_channels = z_channels + # Preserve checkpoint parameter names under flowloss.cfm.model.* while + # keeping the inference-only CFM sampling logic local to FlowLoss. + self.cfm = nn.Module() + self.cfm.add_module("model", DiT(in_channels=z_channels, llm_cond_dim=llm_cond_dim, **kwargs)) + + @torch.no_grad() + def sample(self, z, latent_history, cfg=2.0, patch_size=1, sigma=0.25, temperature=0): + # z: [Batch, Time=1, Dimension] + # latent_history: [Batch, History_Time, Dimension] + if z.ndim != 3: + raise ValueError(f"Expected conditioning rank-3 [Batch, Time, Dimension], got {tuple(z.shape)}") + if latent_history.ndim != 3: + raise ValueError( + f"Expected latent_history rank-3 [Batch, Time, Dimension], got {tuple(latent_history.shape)}" + ) + if z.shape[0] != latent_history.shape[0]: + raise ValueError( + f"Batch mismatch across conditioning and latent_history: {z.shape[0]}, {latent_history.shape[0]}" + ) + dtype = z.dtype + noise = torch.randn(z.shape[0], self.z_channels, patch_size, device=z.device, dtype=torch.float32) + + def fn(t, x): + x = x.to(dtype) + t = t.to(dtype) + if cfg < 1e-5: + if t.ndim == 0: + t = t.repeat(x.shape[0]) + pred = self.cfm.model(x=x, t=t, c=torch.zeros_like(z), latent_history=latent_history) + return pred[:, -patch_size:, :].float() + + pred_cfg = self.cfm.model.forward_with_cfg( + x=x, t=t, c=z, latent_history=latent_history, cfg_scale=cfg, patch_size=patch_size + ) + pred, null_pred = torch.chunk(pred_cfg, 2, dim=0) + return (pred + (pred - null_pred) * cfg).float() + + t = build_timesteps( + 10, + device=noise.device, + dtype=noise.dtype, + use_epss=True, + sway_sampling_coef=-1.0, + ) + return Solver(fn, noise.transpose(1, 2), sigma=sigma, temperature=temperature).integrate(t)[-1].to(dtype) diff --git a/vllm_omni/model_executor/models/ming_tts/fm/__init__.py b/vllm_omni/model_executor/models/ming_tts/fm/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/fm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/ming_tts/fm/dit.py b/vllm_omni/model_executor/models/ming_tts/fm/dit.py new file mode 100644 index 000000000..9aa05d4a2 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/fm/dit.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/dit.py + +import torch +import torch.nn as nn +from x_transformers.x_transformers import RotaryEmbedding + +from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer + + +class DiT(nn.Module): + def __init__( + self, + in_channels=4, + hidden_size=1024, + depth=28, + num_heads=16, + mlp_ratio=4.0, + llm_cond_dim=896, + **kwargs, + ): + super().__init__() + + self.in_channels = in_channels + self.out_channels = in_channels + self.num_heads = num_heads + self.t_embedder = DiTTimestepEmbedding(hidden_size) + self.x_embedder = nn.Linear(in_channels, hidden_size) + self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size) + self.hidden_size = hidden_size + self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) + self.blocks = nn.ModuleList( + [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] + ) + self.final_layer = FinalLayer(hidden_size, self.out_channels) + + def forward(self, x, t, c, latent_history, mask=None): + if x.ndim != 3: + raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}") + if latent_history.ndim != 3: + raise ValueError( + f"Expected latent_history rank-3 [Batch, Time, Dimension], got {tuple(latent_history.shape)}" + ) + if c.ndim != 3: + raise ValueError(f"Expected conditioning rank-3 [Batch, Time, Dimension], got {tuple(c.shape)}") + if x.shape[0] != latent_history.shape[0] or x.shape[0] != c.shape[0]: + raise ValueError( + "Batch mismatch across x, conditioning, and latent_history: " + f"{x.shape[0]}, {c.shape[0]}, {latent_history.shape[0]}" + ) + if x.shape[-1] != self.in_channels: + raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}") + if latent_history.shape[-1] != self.in_channels: + raise ValueError( + f"latent_history feature dim mismatch: got {latent_history.shape[-1]}, expected {self.in_channels}" + ) + if t.ndim == 0: + t = t.reshape(1) + if t.ndim != 1: + raise ValueError(f"Expected timestep rank-1 [Batch], got {tuple(t.shape)}") + if t.shape[0] != x.shape[0]: + raise ValueError(f"Timestep batch mismatch: got {t.shape[0]}, expected {x.shape[0]}") + + t = self.t_embedder(t).unsqueeze(1) + x_now = self.x_embedder(x) + x_history = self.x_embedder(latent_history) + x = torch.cat([x_history, x_now], dim=1) + c = self.c_embedder(c) + y = t + c + x = torch.cat([y, x], dim=1) + rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) + + if mask is not None: + if mask.ndim != 2: + raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}") + if mask.shape[0] != x_now.shape[0] or mask.shape[1] != x_now.shape[1]: + raise ValueError( + f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x_now.shape[0], x_now.shape[1])}" + ) + mask_pad = mask.clone().detach()[:, :1].expand(-1, x_history.shape[1] + c.shape[1]) + mask = torch.cat([mask_pad, mask], dim=-1) + for block in self.blocks: + x = block(x, mask, rope) + x = self.final_layer(x) + return x + + def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size): + if not cfg_scale == 1: + x = torch.cat([x, x], dim=0) + latent_history = torch.cat([latent_history, latent_history], dim=0) + fake_latent = torch.zeros_like(c) + c = torch.cat([c, fake_latent], dim=0) + if t.ndim == 0: + t = t.repeat(x.shape[0]) + model_out = self.forward(x, t, c, latent_history) + return model_out[:, -patch_size:, :] diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts.py b/vllm_omni/model_executor/models/ming_tts/ming_tts.py new file mode 100644 index 000000000..667aded27 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from collections.abc import Iterable +from functools import cached_property +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.model_executor.models import SupportsPP +from vllm.model_executor.models.utils import init_vllm_registered_model +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin + +from .audio_prep import ( + _coerce_prompt_latents, + _find_audio_placeholder_positions, + _find_speaker_placeholder_positions, + _initial_history, + _take_scalar, + coerce_speaker_embeddings, +) +from .config_ming_tts import ( + AUDIO_START_TOKEN_ID, + KEY_CFG, + KEY_DECODE_STEP, + KEY_LAST_STOP_PROB, + KEY_LATENT_HISTORY, + KEY_MAX_DECODE_STEPS, + KEY_MIN_DECODE_STEPS, + KEY_NEXT_EMBEDS, + KEY_PROMPT_LATENT_TAIL, + KEY_REQUEST_ID, + KEY_SIGMA, + KEY_SPEAKER_EMBEDDING, + KEY_TEMPERATURE, + KEY_TEXT_MODE, + MingTTSConfig, +) +from .patch_emission import MING_STOP_REASON_KEY +from .prompt_encoder import _resolve_prompt_latents + + +class _ModelSampleAdapter(nn.Module): + def __init__(self, model: nn.Module): + super().__init__() + self.model = model + + def forward(self, logits, sampling_metadata): + return self.model.sample(logits, sampling_metadata) + + +class MingTTSForConditionalGeneration(nn.Module, SupportsPP, CustomProcessMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + del prefix + self.vllm_config = vllm_config + self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config) + self.ming_config.validate() + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.requires_raw_input_tokens = False + self.model_stage = vllm_config.model_config.model_stage + self._prompt_encoder = None + + if self.model_stage == "llm": + self.model = init_vllm_registered_model(vllm_config=vllm_config, architectures=["MingLLMModel"]) + self.has_preprocess = True + self.has_postprocess = True + self.set_custom_preprocess(self.preprocess) + self.set_custom_postprocess(self.postprocess) + elif self.model_stage == "audio_vae": + self.model = init_vllm_registered_model(vllm_config=vllm_config, architectures=["MingAudioVAEModel"]) + self.requires_raw_input_tokens = True + else: + raise ValueError(f"Invalid Ming model_stage={self.model_stage}") + self.make_empty_intermediate_tensors = getattr(self.model, "make_empty_intermediate_tensors", lambda: None) + + @cached_property + def sampler(self): + if hasattr(self.model, "sample"): + return _ModelSampleAdapter(self.model) + if hasattr(self.model, "sampler"): + return self.model.sampler + return Sampler() + + def embed_input_ids(self, input_ids: torch.Tensor, **kwargs: Any) -> torch.Tensor: + return self.model.embed_input_ids(input_ids=input_ids, **kwargs) + + def forward(self, *args: Any, **kwargs: Any): + return self.model(*args, **kwargs) + + def compute_logits(self, hidden_states, sampling_metadata=None): + return self.model.compute_logits(hidden_states, sampling_metadata=sampling_metadata) + + def sample(self, logits, sampling_metadata): + return self.model.sample(logits, sampling_metadata) if hasattr(self.model, "sample") else None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = list(weights) + if self.model_stage == "llm": + allowed = ("model.", "linear_proj_audio.", "flowloss.", "stop_head.", "spk_head.") + llm_weights = [(k, v) for k, v in weights if k.startswith(allowed)] + if not llm_weights: + raise RuntimeError( + "Ming Stage-0 received no loadable checkpoint weights. " + "Expected prefixes: model.*, linear_proj_audio.*, flowloss.*, stop_head.*, spk_head.*" + ) + loaded = self.model.load_weights(llm_weights) + return {f"model.{name}" for name in loaded} + + audio_weights = [(k, v) for k, v in weights if k.startswith("audio.")] + if not audio_weights: + raise RuntimeError("Ming Stage-1 received no loadable checkpoint weights. Expected prefix: audio.*") + loaded = self.model.load_weights(audio_weights) + return {f"model.{name}" for name in loaded} + + def _resolve_prompt_latents(self, info_dict: dict[str, Any]): + return _resolve_prompt_latents(self, info_dict) + + def preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor | None, **info_dict: Any): + if self.model_stage != "llm": + return input_ids, input_embeds, {} + input_embeds = self.model.embed_input_ids(input_ids).clone() + return ( + self._prefill_preprocess(input_ids, input_embeds, **info_dict) + if int(input_ids.shape[0]) > 1 + else self._decode_preprocess(input_ids, input_embeds, **info_dict) + ) + + def postprocess(self, hidden_states: torch.Tensor, **info_dict: Any) -> dict[str, Any]: + if self.model_stage != "llm" or hidden_states.numel() == 0: + return {} + req_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + pending = self.model.pop_postprocess_update(req_id) + if not pending or not isinstance(pending.get("ming_latent_patch"), torch.Tensor): + return {} + + update = { + KEY_LATENT_HISTORY: pending[KEY_LATENT_HISTORY].detach().to("cpu").contiguous(), + KEY_NEXT_EMBEDS: pending[KEY_NEXT_EMBEDS].detach().to("cpu").contiguous(), + KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0)) + 1, + } + stop_prob = _take_scalar(pending.get("ming_stop_prob"), 0) + if stop_prob is not None: + update[KEY_LAST_STOP_PROB] = stop_prob + stop_reason = pending.get(MING_STOP_REASON_KEY) + if isinstance(stop_reason, str): + update[MING_STOP_REASON_KEY] = stop_reason + if isinstance(req_id, str): + update[KEY_REQUEST_ID] = req_id + return update + + def _prefill_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: Any): + if bool(info_dict.get(KEY_TEXT_MODE, False)): + update: dict[str, Any] = {KEY_TEXT_MODE: True} + request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + if request_id is not None: + update[KEY_REQUEST_ID] = request_id + if int(input_ids.shape[0]) > 1 and int(input_ids[-1].item()) == AUDIO_START_TOKEN_ID: + return input_ids[:-1], input_embeds[:-1], update + return input_ids, input_embeds, update + + update: dict[str, Any] = {KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0))} + prompt_latents = _resolve_prompt_latents(self, info_dict) + history = _initial_history( + prompt_latents["frames"] if prompt_latents is not None else None, + history_size=self.ming_config.history_patch_size, + latent_dim=self.ming_config.latent_dim, + device=input_embeds.device, + dtype=torch.float32, + ) + update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous() + update[KEY_PROMPT_LATENT_TAIL] = update[KEY_LATENT_HISTORY] + + speaker_embedding = info_dict.get(KEY_SPEAKER_EMBEDDING, info_dict.get("speaker_embedding")) + speaker_embeddings = None + if speaker_embedding is not None: + speaker_embeddings = coerce_speaker_embeddings( + speaker_embedding, + use_zero_spk_emb=bool(info_dict.get("use_zero_spk_emb", False)), + ) + if speaker_embeddings is not None and len(speaker_embeddings) > 0: + speaker_slots = _find_speaker_placeholder_positions(input_ids, self.vllm_config.model_config.hf_config) + if len(speaker_slots) < len(speaker_embeddings): + raise RuntimeError( + "Could not locate enough speaker placeholder slots: " + f"found {len(speaker_slots)}, need {len(speaker_embeddings)}" + ) + for speaker_slot, spk in zip(speaker_slots, speaker_embeddings): + spk_proj = self.model.project_speaker_embedding( + spk.to(device=input_embeds.device, dtype=input_embeds.dtype).unsqueeze(0) + ).squeeze(0) + input_embeds[speaker_slot] = spk_proj + + if prompt_latents is not None and prompt_latents["patches"] is not None: + prompt_patches = prompt_latents["patches"].to(dtype=getattr(self.model, "fm_dtype", torch.float32)) + prompt_embeds = self.model.linear_proj_audio(prompt_patches).squeeze(1) + placeholder_pos = _find_audio_placeholder_positions(input_ids, self.ming_config) + take = min(int(placeholder_pos.numel()), int(prompt_embeds.shape[0])) + if take > 0: + input_embeds[placeholder_pos[:take]] = prompt_embeds[:take].to(dtype=input_embeds.dtype) + + request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + if request_id is not None: + update[KEY_REQUEST_ID] = request_id + _copy_runtime_controls(update, info_dict) + return input_ids, input_embeds, update + + def _decode_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: Any): + if bool(info_dict.get(KEY_TEXT_MODE, False)): + update: dict[str, Any] = {KEY_TEXT_MODE: True} + request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + if request_id is not None: + update[KEY_REQUEST_ID] = request_id + return input_ids, input_embeds, update + + update: dict[str, Any] = {KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0))} + history = info_dict.get(KEY_LATENT_HISTORY) + if isinstance(history, torch.Tensor): + update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous() + else: + zero_history = torch.zeros( + (self.ming_config.history_patch_size, self.ming_config.latent_dim), + device=input_embeds.device, + dtype=torch.float32, + ) + update[KEY_LATENT_HISTORY] = zero_history.detach().to("cpu").contiguous() + + next_embeds = info_dict.get(KEY_NEXT_EMBEDS) + if isinstance(next_embeds, torch.Tensor) and input_ids.numel() == 1: + if not torch.isfinite(next_embeds).all(): + raise RuntimeError("Non-finite next_embeds before decode preprocess write.") + input_embeds[0] = ( + next_embeds.detach() + .reshape(-1, self.ming_config.llm_hidden_size)[0] + .to( + device=input_embeds.device, + dtype=input_embeds.dtype, + ) + ) + if not torch.isfinite(input_embeds[0]).all(): + raise RuntimeError("Non-finite backbone input_embeds after decode preprocess write.") + + request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + if request_id is not None: + update[KEY_REQUEST_ID] = request_id + _copy_runtime_controls(update, info_dict) + return input_ids, input_embeds, update + + +def _copy_runtime_controls(update: dict[str, Any], info_dict: dict[str, Any]) -> None: + for key in (KEY_CFG, KEY_SIGMA, KEY_TEMPERATURE, KEY_MAX_DECODE_STEPS, KEY_MIN_DECODE_STEPS): + if key in info_dict: + update[key] = info_dict[key] + + +__all__ = [ + "MingTTSForConditionalGeneration", + "_ModelSampleAdapter", + "_coerce_prompt_latents", + "_find_audio_placeholder_positions", + "_find_speaker_placeholder_positions", + "_initial_history", +] diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_audio_vae.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_audio_vae.py new file mode 100644 index 000000000..8b029ffba --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_audio_vae.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/modeling_audio_vae.py +from __future__ import annotations + +import os +import time +import warnings +from collections import OrderedDict +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .config_ming_tts import KEY_CHUNK_ID, KEY_REQUEST_ID, MingTTSConfig +from .patch_emission import MING_STOP_REASON_KEY + +logger = init_logger(__name__) + +MING_FINAL_DECODE_STEP_KEY = "ming_final_decode_step" +_STREAM_STATE_TTL_SECONDS = 15 * 60 +_MAX_STREAM_STATES = 1024 + + +class MingAudioVAEModel(nn.Module): + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + del prefix + self.vllm_config = vllm_config + self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config) + if self.ming_config.audio_tokenizer_config is None: + raise ValueError("MingAudioVAEModel requires audio_tokenizer_config") + self.ming_config.validate() + + self.audio = AudioVAE(self.ming_config.audio_tokenizer_config) + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + self._past_key_values: dict[str, Any] = {} + self._stream_state: dict[str, tuple[Any, Any, Any]] = {} + self._patch_totals: dict[str, int] = {} + self._sample_totals: dict[str, int] = {} + self._state_access_times: OrderedDict[str, float] = OrderedDict() + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + hidden_size = int(self.ming_config.llm_hidden_size) + if input_ids is None or input_ids.numel() == 0: + return torch.empty((0, hidden_size), device=self.vllm_config.device_config.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], hidden_size), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None: + del hidden_states, sampling_metadata + return None + + def chunked_decode_streaming( + self, + latent_chunk: torch.Tensor, + *, + request_id: str, + finished: bool, + ) -> tuple[torch.Tensor, Any, Any, bool, bool]: + had_past_key_values = request_id in self._past_key_values + had_stream_state = _has_stream_state(self._stream_state.get(request_id)) + stream_state = self._stream_state.get(request_id, (None, None, None)) + past_key_values = self._past_key_values.get(request_id) + waveform_parts: list[torch.Tensor] = [] + + patch_count = int(latent_chunk.shape[0]) + for patch_idx in range(patch_count): + # [Batch, Time, Dimension] = [1, patch_size, latent_dim] + latent_patch = latent_chunk[patch_idx : patch_idx + 1] + is_last_patch = finished and patch_idx == patch_count - 1 + waveform, stream_state, past_key_values = self.audio.decode( + latent_patch, + past_key_values=past_key_values, + use_cache=True, + stream_state=stream_state, + last_chunk=is_last_patch, + ) + waveform_parts.append(waveform.reshape(-1).to(torch.float32)) + + waveform_flat = ( + torch.cat(waveform_parts, dim=0) + if waveform_parts + else torch.zeros((0,), dtype=torch.float32, device=latent_chunk.device) + ) + return waveform_flat, stream_state, past_key_values, had_past_key_values, had_stream_state + + @torch.inference_mode() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + model_intermediate_buffer: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + del input_ids, positions, intermediate_tensors, inputs_embeds + + info_list = model_intermediate_buffer if isinstance(model_intermediate_buffer, list) else None + if info_list is None: + runtime_infos = kwargs.get("runtime_additional_information") + info_list = runtime_infos if isinstance(runtime_infos, list) else None + info_list = info_list or [{}] + num_reqs = max(len(info_list), 1) + sr_tensor = torch.tensor(int(self.ming_config.sample_rate), dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32, device=self.vllm_config.device_config.device) + + outputs: list[torch.Tensor] = [] + srs: list[torch.Tensor] = [] + + for idx in range(num_reqs): + info = info_list[idx] if idx < len(info_list) and isinstance(info_list[idx], dict) else {} + meta = info.get("meta") or {} + kv_meta = info.get("kv_metadata") or {} + has_ming_context = _has_ming_context(info) + if has_ming_context and KEY_REQUEST_ID not in info and KEY_REQUEST_ID not in kv_meta: + raise RuntimeError( + f"Ming Stage-2 received a payload without {KEY_REQUEST_ID}. keys={sorted(info.keys())}" + ) + request_id = _resolve_request_id(info, idx) + self._touch_request_state(request_id) + chunk_id = _coerce_optional_int(_first_present(info.get(KEY_CHUNK_ID), kv_meta.get(KEY_CHUNK_ID))) + finished = _coerce_finished( + _first_present( + meta.get("stream_finished"), + meta.get("finished"), + info.get("stream_finished"), + kv_meta.get("stream_finished"), + torch.tensor(True), + ) + ) + latent = _first_present( + info.get("latent"), + info.get("ming_latent_patches"), + kv_meta.get("ming_latent_patches"), + ) + stripped = bool( + _first_present( + meta.get("_ming_payload_stripped"), + info.get("_ming_payload_stripped"), + kv_meta.get("_ming_payload_stripped"), + False, + ) + ) + if stripped: + raise RuntimeError( + "Ming Stage-2 payload was stripped before model entry. " + f"request_id={request_id} chunk_id={chunk_id} keys={sorted(info.keys())}" + ) + if latent is None: + if has_ming_context and not finished: + raise RuntimeError( + "Ming Stage-2 received no latent chunk for an unfinished request. " + f"request_id={request_id} chunk_id={chunk_id} keys={sorted(info.keys())}" + ) + if finished: + self._clear_request_state(request_id) + outputs.append(empty) + srs.append(sr_tensor) + continue + + latent_tensor = _coerce_latent_chunk( + latent, + device=self.vllm_config.device_config.device, + dtype=next(self.audio.parameters()).dtype, + latent_dim=self.ming_config.latent_dim, + patch_size=self.ming_config.patch_size, + ) + if latent_tensor is None or latent_tensor.numel() == 0: + if not finished: + raise RuntimeError( + "Ming Stage-2 received an empty latent chunk before final flush. " + f"request_id={request_id} chunk_id={chunk_id} latent_shape={_shape_of(latent_tensor)}" + ) + if finished: + self._clear_request_state(request_id) + outputs.append(empty) + srs.append(sr_tensor) + continue + + patch_count = int(latent_tensor.shape[0]) + waveform_flat, stream_state, past_key_values, had_past_key_values, had_stream_state = ( + self.chunked_decode_streaming( + latent_tensor, + request_id=request_id, + finished=finished, + ) + ) + total_patch_count = self._patch_totals.get(request_id, 0) + patch_count + total_waveform_numel = self._sample_totals.get(request_id, 0) + int(waveform_flat.numel()) + self._patch_totals[request_id] = total_patch_count + self._sample_totals[request_id] = total_waveform_numel + if (had_past_key_values or had_stream_state) and waveform_flat.numel() == 0 and not finished: + raise RuntimeError( + "Ming Stage-2 produced an empty waveform after cached streaming state already existed. " + f"request_id={request_id} chunk_id={chunk_id} latent_shape={tuple(latent_tensor.shape)} " + f"had_past_key_values={had_past_key_values} had_stream_state={had_stream_state}" + ) + + if finished: + if os.environ.get("MING_TTS_STAGE1_FINAL_LOG") == "1": + logger.info( + "MING_STAGE1_FINAL %s", + { + "request_id": request_id, + "chunk_id": chunk_id, + "stop_reason": _first_present( + info.get(MING_STOP_REASON_KEY), + kv_meta.get(MING_STOP_REASON_KEY), + ), + "final_decode_step": _coerce_optional_int( + _first_present( + info.get(MING_FINAL_DECODE_STEP_KEY), + kv_meta.get(MING_FINAL_DECODE_STEP_KEY), + ) + ), + "final_chunk_patch_count": patch_count, + "total_patch_count": total_patch_count, + "final_chunk_waveform_numel": int(waveform_flat.numel()), + "total_waveform_numel": total_waveform_numel, + }, + ) + self._clear_request_state(request_id) + else: + self._past_key_values[request_id] = past_key_values + self._stream_state[request_id] = stream_state + + outputs.append(waveform_flat) + srs.append(sr_tensor) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": outputs, + "sr": srs, + }, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = list(weights) + if not weights: + raise RuntimeError("MingAudioVAEModel received no checkpoint weights.") + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded: set[str] = set() + skipped: list[str] = [] + + for name, loaded_weight in weights: + if name not in params_dict: + skipped.append(name) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded.add(name) + + missing = {name for name in params_dict if name.startswith("audio.")} - loaded + if missing: + raise RuntimeError(f"MingAudioVAEModel: {len(missing)} params not loaded. First few: {sorted(missing)[:5]}") + if skipped: + warnings.warn( + f"MingAudioVAEModel: skipped {len(skipped)} checkpoint keys during load. First few: {skipped[:8]}", + stacklevel=2, + ) + return loaded + + def _clear_request_state(self, request_id: str) -> None: + self._past_key_values.pop(request_id, None) + self._stream_state.pop(request_id, None) + self._patch_totals.pop(request_id, None) + self._sample_totals.pop(request_id, None) + self._state_access_times.pop(request_id, None) + + def _touch_request_state(self, request_id: str) -> None: + now = time.monotonic() + self._evict_expired_request_states(now) + self._state_access_times[request_id] = now + self._state_access_times.move_to_end(request_id) + while len(self._state_access_times) > _MAX_STREAM_STATES: + evicted_request_id, _ = self._state_access_times.popitem(last=False) + self._clear_request_state(evicted_request_id) + + def _evict_expired_request_states(self, now: float) -> None: + cutoff = now - _STREAM_STATE_TTL_SECONDS + while self._state_access_times: + request_id, last_access = next(iter(self._state_access_times.items())) + if last_access >= cutoff: + break + self._clear_request_state(request_id) + + +def _coerce_finished(value: Any) -> bool: + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return False + return bool(value.reshape(-1)[0].item()) + return bool(value) + + +def _first_present(*values: Any) -> Any: + for value in values: + if value is not None: + return value + return None + + +def _coerce_optional_int(value: Any, default: int | None = None) -> int | None: + if value is None: + return default + if isinstance(value, list): + value = value[0] if value else default + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return default + return int(value.reshape(-1)[0].item()) + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _resolve_request_id(info: dict[str, Any], idx: int) -> str: + kv_meta = info.get("kv_metadata") or {} + request_id = info.get(KEY_REQUEST_ID) or kv_meta.get(KEY_REQUEST_ID) + if request_id is None: + return str(idx) + if not isinstance(request_id, str) or not request_id: + raise RuntimeError(f"Ming Stage-2 received invalid request id: {request_id!r}") + return request_id + + +def _has_ming_context(info: dict[str, Any]) -> bool: + kv_meta = info.get("kv_metadata") or {} + keys = ( + KEY_REQUEST_ID, + KEY_CHUNK_ID, + "ming_latent_patches", + "_ming_payload_stripped", + ) + return any(key in info or key in kv_meta for key in keys) or "latent" in info + + +def _shape_of(value: Any) -> tuple[int, ...] | None: + if isinstance(value, torch.Tensor): + return tuple(value.shape) + return None + + +def _has_stream_state(value: Any) -> bool: + if not isinstance(value, tuple): + return False + return any(item is not None for item in value) + + +def _coerce_latent_chunk( + value: Any, + *, + device: torch.device, + dtype: torch.dtype, + latent_dim: int, + patch_size: int, +) -> torch.Tensor | None: + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + + latent = value.detach() + if latent.ndim == 2: + latent = latent.unsqueeze(0) + if latent.ndim != 3: + raise ValueError(f"Expected latent chunk rank-3 [Batch, Time, Dimension], got {tuple(latent.shape)}") + if latent.shape[-2] != patch_size: + raise ValueError(f"Latent patch size mismatch: got {latent.shape[-2]}, expected {patch_size}") + if latent.shape[-1] != latent_dim: + raise ValueError(f"Latent dim mismatch: got {latent.shape[-1]}, expected {latent_dim}") + return latent.to(device=device, dtype=dtype) diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py new file mode 100644 index 000000000..92daad967 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py +from __future__ import annotations + +import warnings +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.qwen2 import Qwen2Model +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix +from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .aggregator import Aggregator +from .config_ming_tts import ( + KEY_CFG, + KEY_DECODE_STEP, + KEY_LATENT_HISTORY, + KEY_MAX_DECODE_STEPS, + KEY_MIN_DECODE_STEPS, + KEY_NEXT_EMBEDS, + KEY_REQUEST_ID, + KEY_SIGMA, + KEY_TEMPERATURE, + KEY_TEXT_MODE, + MingTTSConfig, +) +from .flowloss_head import FlowLoss +from .patch_emission import ( + MING_STOP_REASON_CODES, + MING_STOP_REASON_CONTINUE, + MING_STOP_REASON_KEY, + MING_STOP_REASON_MAX_DECODE_STEPS, + MING_STOP_REASON_STOP_HEAD, + _coerce_latent_history, + _get_request_token_counts, + _normalize_request_infos, + _resolve_ming_stop_decision, + _resolve_optional_runtime_int, + _resolve_runtime_float, + _resolve_runtime_int, + _validate_ming_decode_window, +) + +logger = init_logger(__name__) + + +class MingLLMModel(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.model.": "model.", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config) + self.ming_config.validate() + self.vllm_config = vllm_config + self.prefix = prefix + self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config) + self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + self.linear_proj_audio = Aggregator( + in_channels=self.ming_config.latent_dim, + llm_input_dim=self.ming_config.llm_hidden_size, + **self.ming_config.aggregator_config, + ) + self.flowloss = FlowLoss( + z_channels=self.ming_config.latent_dim, + llm_cond_dim=self.ming_config.llm_hidden_size, + **self.ming_config.ditar_config, + ) + self.stop_head = nn.Linear(self.ming_config.llm_hidden_size, 2, bias=True) + self.spk_head = nn.Linear(192, self.ming_config.llm_hidden_size, bias=True) + self.flowloss.to(dtype=self.fm_dtype) + self.linear_proj_audio.to(dtype=self.fm_dtype) + self.stop_head.to(dtype=self.fm_dtype) + self.spk_head.to(dtype=self.fm_dtype) + self._pending_postprocess_updates: dict[str, dict[str, Any]] = {} + self._last_ming_next_token_ids = None + self._last_text_mode = False + + def embed_input_ids( + self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, **_: Any + ) -> torch.Tensor: + if inputs_embeds is not None: + return inputs_embeds + return self.model.embed_input_ids(input_ids) + + def project_speaker_embedding(self, spk_emb: torch.Tensor) -> torch.Tensor: + return self.spk_head(spk_emb) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + latent_history: torch.Tensor | None = None, + model_intermediate_buffer: list[dict[str, Any]] | None = None, + seq_token_counts: list[int] | None = None, + **kwargs: object, + ) -> OmniOutput | IntermediateTensors | torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + if model_intermediate_buffer is None: + model_intermediate_buffer = kwargs.get("runtime_additional_information") + request_infos = _normalize_request_infos(model_intermediate_buffer) + _validate_ming_decode_window( + request_infos, + min_stop_step=int(self.ming_config.stop_head_min_steps), + default_max_decode_steps=self.ming_config.max_decode_steps, + ) + backbone_out = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + if isinstance(backbone_out, IntermediateTensors): + self._last_ming_next_token_ids = None + return backbone_out + hidden_states = _extract_hidden_states(backbone_out) + token_counts = _get_request_token_counts(hidden_states, request_infos, seq_token_counts) + text_mode = bool(request_infos) and all(bool(info.get(KEY_TEXT_MODE, False)) for info in request_infos) + if request_infos and any(bool(info.get(KEY_TEXT_MODE, False)) for info in request_infos) and not text_mode: + raise RuntimeError("Mixed Ming text/audio modes in one Stage-0 batch are unsupported.") + if text_mode: + self._last_ming_next_token_ids = None + self._last_text_mode = True + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={KEY_TEXT_MODE: True}, + intermediate_tensors=intermediate_tensors, + ) + self._last_text_mode = False + if latent_history is None and not token_counts: + self._last_ming_next_token_ids = None + return OmniOutput( + text_hidden_states=hidden_states, multimodal_outputs=None, intermediate_tensors=intermediate_tensors + ) + if latent_history is not None and not token_counts: + token_counts, request_infos = [hidden_states.shape[0]], [{KEY_LATENT_HISTORY: latent_history}] + + total_tokens = hidden_states.shape[0] + latent_patch_tokens = next_embed_tokens = new_history_tokens = None + decode_step_tokens = has_patch = stop_reason_code_tokens = None + next_token_ids = [] + pending_updates: dict[str, dict[str, Any]] = {} + cursor = 0 + any_decode = False + for req_idx, token_count in enumerate(token_counts): + end = min(cursor + token_count, total_tokens) + if end <= cursor: + continue + req_info = request_infos[req_idx] if req_idx < len(request_infos) else {} + req_id = req_info.get(KEY_REQUEST_ID) + req_history = _coerce_latent_history( + req_info.get(KEY_LATENT_HISTORY), device=hidden_states.device, dtype=self.fm_dtype, cfg=self.ming_config + ) + if req_history is None: + cursor = end + continue + decode_step = int(req_info.get(KEY_DECODE_STEP, req_info.get("generated_len", 0))) + decode_hidden, output_index = ( + (hidden_states[cursor:end], cursor) if token_count == 1 else (hidden_states[end - 1 : end], end - 1) + ) + sampled_token_latent, next_embeds, new_history, stop_probs = self._decode_one_step( + hidden_states=decode_hidden, + latent_history=req_history, + cfg_scale=_resolve_runtime_float(req_info, KEY_CFG, self.ming_config.cfg), + sigma=_resolve_runtime_float(req_info, KEY_SIGMA, self.ming_config.sigma), + temperature=_resolve_runtime_float(req_info, KEY_TEMPERATURE, self.ming_config.temperature), + ) + req_max_decode_steps = _resolve_runtime_int( + req_info, KEY_MAX_DECODE_STEPS, self.ming_config.max_decode_steps + ) + req_min_decode_steps = _resolve_optional_runtime_int(req_info, KEY_MIN_DECODE_STEPS, 0) + if latent_patch_tokens is None: + latent_patch_tokens = sampled_token_latent.new_zeros( + (total_tokens, self.ming_config.patch_size, self.ming_config.latent_dim) + ) + next_embed_tokens = next_embeds.new_zeros((total_tokens, 1, self.ming_config.llm_hidden_size)) + new_history_tokens = new_history.new_zeros( + (total_tokens, self.ming_config.history_patch_size, self.ming_config.latent_dim) + ) + decode_step_tokens = torch.zeros((total_tokens,), dtype=torch.int32, device=hidden_states.device) + has_patch = torch.zeros((total_tokens,), dtype=torch.bool, device=hidden_states.device) + stop_reason_code_tokens = torch.zeros((total_tokens,), dtype=torch.int32, device=hidden_states.device) + latent_patch_tokens[output_index : output_index + 1] = sampled_token_latent + next_embed_tokens[output_index : output_index + 1] = next_embeds + new_history_tokens[output_index : output_index + 1] = new_history + decode_step_tokens[output_index : output_index + 1] = decode_step + has_patch[output_index : output_index + 1] = True + stop_reason, _, _, _, next_token_id = _resolve_ming_stop_decision( + step=decode_step, + stop_prob=float(stop_probs.reshape(-1)[0].item()), + stop_threshold=float(self.ming_config.stop_head_threshold), + min_stop_step=int(self.ming_config.stop_head_min_steps), + min_decode_steps=req_min_decode_steps, + max_decode_steps=req_max_decode_steps, + audio_dummy_token_id=int(self.ming_config.audio_dummy_token_id), + text_eos_token_id=int(self.ming_config.text_eos_token_id), + ) + stop_reason_code_tokens[output_index : output_index + 1] = MING_STOP_REASON_CODES[stop_reason] + next_token_ids.append(int(next_token_id)) + if isinstance(req_id, str): + pending_updates[req_id] = { + KEY_LATENT_HISTORY: new_history, + KEY_NEXT_EMBEDS: next_embeds, + "ming_latent_patch": sampled_token_latent, + "ming_stop_prob": stop_probs, + MING_STOP_REASON_KEY: stop_reason, + } + any_decode = True + cursor = end + + self._pending_postprocess_updates = pending_updates + if not any_decode: + self._last_ming_next_token_ids = None + return OmniOutput( + text_hidden_states=hidden_states, multimodal_outputs=None, intermediate_tensors=intermediate_tensors + ) + self._last_ming_next_token_ids = next_token_ids + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={ + "ming_latent_patch": latent_patch_tokens, + "ming_next_embeds": next_embed_tokens, + "ming_new_history": new_history_tokens, + "ming_decode_step": decode_step_tokens, + "ming_has_patch": has_patch, + MING_STOP_REASON_KEY: stop_reason_code_tokens, + }, + intermediate_tensors=intermediate_tensors, + ) + + def pop_postprocess_update(self, req_id: str) -> dict[str, Any]: + return self._pending_postprocess_updates.pop(req_id, {}) if isinstance(req_id, str) else {} + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata) -> torch.Tensor | None: + del sampling_metadata + if self._last_text_mode: + return ( + None + if hidden_states is None or hidden_states.numel() == 0 + else self.model.compute_logits(hidden_states) + ) + if hidden_states is None or hidden_states.numel() == 0: + return None + if hidden_states.dim() != 2: + raise RuntimeError( + f"Expected hidden_states rank-2 [B,H] in compute_logits, got {tuple(hidden_states.shape)}" + ) + batch_size = hidden_states.shape[0] + next_token_ids = self._last_ming_next_token_ids + self._last_ming_next_token_ids = None + if next_token_ids is None: + logger.debug( + "Missing Ming forced next-token ids before compute_logits. " + "Using dummy next-token IDs (likely during a profiling or dummy run)." + ) + next_token_ids = [int(self.ming_config.text_eos_token_id)] * batch_size + if len(next_token_ids) != batch_size: + raise RuntimeError( + "Ming forced next-token batch mismatch: " + f"got {len(next_token_ids)} ids for hidden_states batch {batch_size}." + ) + logits = torch.full( + (batch_size, self.ming_config.llm_vocab_size), + float("-inf"), + device=hidden_states.device, + dtype=torch.float32, + ) + for i, next_token_id in enumerate(next_token_ids): + logits[i, int(next_token_id)] = 0.0 + return logits + + def sample(self, logits, sampling_metadata): + if logits is None: + return None + if self._last_text_mode: + return self.model.sample(logits, sampling_metadata) + del sampling_metadata + return SamplerOutput( + sampled_token_ids=logits.argmax(dim=-1, keepdim=True).to(dtype=torch.int32), + logprobs_tensors=None, + ) + + def _decode_one_step( + self, + *, + hidden_states: torch.Tensor, + latent_history: torch.Tensor, + cfg_scale: float, + sigma: float, + temperature: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if hidden_states.dim() != 2: + raise RuntimeError(f"Expected decode hidden_states rank-2 [B,H], got {tuple(hidden_states.shape)}") + if latent_history.dim() != 3: + raise RuntimeError(f"Expected latent_history rank-3 [B,T,D], got {tuple(latent_history.shape)}") + if hidden_states.shape[0] != latent_history.shape[0]: + raise RuntimeError( + "Batch mismatch: " + f"hidden_states B={hidden_states.shape[0]} " + f"vs latent_history B={latent_history.shape[0]}" + ) + # [Batch, Hidden] -> [Batch, Time, Hidden] = [B, 1, H] for FlowLoss conditioning. + z_diff_cond = hidden_states.to(dtype=self.fm_dtype).unsqueeze(1) + if not torch.isfinite(z_diff_cond).all(): + raise RuntimeError("Non-finite z_diff_cond before FlowLoss.sample().") + sampled_token_latent = self.flowloss.sample( + z=z_diff_cond, + latent_history=latent_history, + cfg=cfg_scale, + patch_size=self.ming_config.patch_size, + sigma=sigma, + temperature=temperature, + ) + expected_shape = (hidden_states.shape[0], self.ming_config.patch_size, self.ming_config.latent_dim) + if tuple(sampled_token_latent.shape) != expected_shape: + raise RuntimeError( + f"FlowLoss output shape mismatch: got {tuple(sampled_token_latent.shape)}, expected {expected_shape}" + ) + new_history = torch.cat([latent_history[:, self.ming_config.patch_size :, :], sampled_token_latent], dim=1) + # Aggregator expects [Batch, Time, Dimension] = [B, 4, 64] and returns [B, 1, H]. + next_embeds = self.linear_proj_audio(sampled_token_latent) + stop_probs = self.stop_head(hidden_states.to(dtype=self.fm_dtype)).softmax(dim=-1)[:, 1] + if not torch.isfinite(sampled_token_latent).all(): + raise RuntimeError("Non-finite sampled_token_latent in Ming decode step.") + if not torch.isfinite(next_embeds).all(): + raise RuntimeError("Non-finite next_embeds in Ming decode step.") + if not torch.isfinite(stop_probs).all(): + raise RuntimeError("Non-finite stop_probs in Ming decode step.") + return sampled_token_latent, next_embeds, new_history, stop_probs + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params = AutoWeightsLoader(self).load_weights(weights, mapper=self.hf_to_vllm_mapper) + _warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.", fatal=True) + _warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.", fatal=True) + _warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.", fatal=True) + _warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.", fatal=True) + return loaded_params + + +def _extract_hidden_states(backbone_out: object) -> torch.Tensor: + if isinstance(backbone_out, torch.Tensor): + return backbone_out + if hasattr(backbone_out, "last_hidden_state"): + return backbone_out.last_hidden_state + if isinstance(backbone_out, (tuple, list)) and backbone_out and isinstance(backbone_out[0], torch.Tensor): + return backbone_out[0] + raise TypeError(f"Unsupported backbone forward output type: {type(backbone_out)}") + + +def _resolve_ming_runtime_dtype(vllm_config: VllmConfig) -> torch.dtype: + dtype = getattr(vllm_config.model_config, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str): + normalized = dtype.strip().lower() + if normalized in ("float16", "half", "torch.float16"): + return torch.float16 + if normalized in ("bfloat16", "bf16", "torch.bfloat16"): + return torch.bfloat16 + if normalized in ("float32", "fp32", "torch.float32"): + return torch.float32 + return torch.float32 + + +def _warn_missing_prefix( + module_name: str, + params_dict: dict[str, nn.Parameter], + loaded_params: set[str], + prefix: str, + fatal: bool = False, +) -> None: + missing = {key for key in params_dict if key.startswith(prefix)} - loaded_params + if not missing: + return + msg = ( + f"MingLLMModel: {len(missing)} {module_name} params not loaded " + f"(prefix={prefix}). First few: {sorted(missing)[:5]}" + ) + if fatal: + raise RuntimeError(msg) + warnings.warn(msg, stacklevel=3) + + +__all__ = [ + "Aggregator", + "FlowLoss", + "MING_STOP_REASON_CODES", + "MING_STOP_REASON_CONTINUE", + "MING_STOP_REASON_KEY", + "MING_STOP_REASON_MAX_DECODE_STEPS", + "MING_STOP_REASON_STOP_HEAD", + "MingLLMModel", + "_coerce_latent_history", + "_extract_hidden_states", + "_resolve_ming_runtime_dtype", + "_resolve_ming_stop_decision", + "_warn_missing_prefix", +] diff --git a/vllm_omni/model_executor/models/ming_tts/patch_emission.py b/vllm_omni/model_executor/models/ming_tts/patch_emission.py new file mode 100644 index 000000000..c0a25a1c1 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/patch_emission.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py +from __future__ import annotations + +from typing import Any + +import torch +from vllm.forward_context import get_forward_context, is_forward_context_available + +from .config_ming_tts import KEY_MAX_DECODE_STEPS, KEY_MIN_DECODE_STEPS, KEY_REQUEST_ID, KEY_TEXT_MODE, MingTTSConfig + +MING_STOP_REASON_CONTINUE = "continue" +MING_STOP_REASON_STOP_HEAD = "stop_head" +MING_STOP_REASON_MAX_DECODE_STEPS = "max_decode_steps" +MING_STOP_REASON_KEY = "ming_stop_reason" +MING_STOP_REASON_CODES = { + MING_STOP_REASON_CONTINUE: 0, + MING_STOP_REASON_STOP_HEAD: 1, + MING_STOP_REASON_MAX_DECODE_STEPS: 2, +} + + +def _normalize_request_infos(model_intermediate_buffer: object) -> list[dict[str, Any]]: + if not isinstance(model_intermediate_buffer, list): + return [] + infos: list[dict[str, Any]] = [] + for item in model_intermediate_buffer: + infos.append(item if isinstance(item, dict) else {}) + return infos + + +def _get_request_token_counts( + hidden_states: torch.Tensor, + request_infos: list[dict[str, Any]], + seq_token_counts: list[int] | None, +) -> list[int]: + if seq_token_counts: + return [int(x) for x in seq_token_counts] + + if is_forward_context_available(): + slices = getattr(get_forward_context(), "ubatch_slices", None) + if slices is not None and len(slices) > 0: + counts: list[int] = [] + for item in slices: + if isinstance(item, int): + counts.append(int(item)) + elif hasattr(item, "stop") and hasattr(item, "start"): + counts.append(int(item.stop) - int(item.start)) + if counts: + return counts + + if request_infos: + if len(request_infos) == hidden_states.shape[0]: + return [1] * hidden_states.shape[0] + return [hidden_states.shape[0]] + + return [] + + +def _coerce_latent_history( + value: object, + *, + device: torch.device, + dtype: torch.dtype, + cfg: MingTTSConfig, +) -> torch.Tensor | None: + if value is None: + return None + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + + history = value.detach() + if history.ndim == 2: + history = history.unsqueeze(0) + if history.ndim != 3: + raise RuntimeError(f"Expected latent_history rank-3 [B,T,D], got {tuple(history.shape)}") + if history.shape[1] != cfg.history_patch_size or history.shape[2] != cfg.latent_dim: + raise RuntimeError( + f"latent_history shape mismatch: got {tuple(history.shape)}, " + f"expected [B,{cfg.history_patch_size},{cfg.latent_dim}]" + ) + return history.to(device=device, dtype=dtype) + + +def _resolve_runtime_float(req_info: dict[str, Any], key: str, default_value: float) -> float: + raw = req_info.get(key, default_value) + try: + value = float(raw) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"Invalid {key}: expected float-like value, got {raw!r}") from exc + if not value >= 0.0: + raise RuntimeError(f"Invalid {key}: expected non-negative value, got {value}") + return value + + +def _resolve_runtime_int(req_info: dict[str, Any], key: str, default_value: int) -> int: + raw = req_info.get(key, default_value) + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"Invalid {key}: expected int-like value, got {raw!r}") from exc + if value <= 0: + raise RuntimeError(f"Invalid {key}: expected positive value, got {value}") + return value + + +def _resolve_optional_runtime_int(req_info: dict[str, Any], key: str, default_value: int) -> int: + raw = req_info.get(key, default_value) + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"Invalid {key}: expected int-like value, got {raw!r}") from exc + if value < 0: + raise RuntimeError(f"Invalid {key}: expected non-negative value, got {value}") + return value + + +def _validate_ming_decode_window( + request_infos: list[dict[str, Any]], + *, + min_stop_step: int, + default_max_decode_steps: int, +) -> None: + for i, info in enumerate(request_infos): + if info.get(KEY_TEXT_MODE): + continue + max_steps = _resolve_runtime_int(info, KEY_MAX_DECODE_STEPS, default_max_decode_steps) + min_steps = _resolve_optional_runtime_int(info, KEY_MIN_DECODE_STEPS, 0) + min_required = max(min_stop_step + 1, min_steps) + if max_steps < min_required: + req_id = info.get(KEY_REQUEST_ID, f"idx={i}") + raise ValueError( + f"Ming request {req_id!r}: max_decode_steps={max_steps} < " + f"min_required_decode_steps={min_required} " + f"(min_stop_step={min_stop_step}, min_decode_steps={min_steps})" + ) + + +def _resolve_ming_stop_decision( + *, + step: int, + stop_prob: float, + stop_threshold: float, + min_stop_step: int, + min_decode_steps: int, + max_decode_steps: int, + audio_dummy_token_id: int, + text_eos_token_id: int, +) -> tuple[str, bool, bool, int, int]: + min_required_decode_steps = max(min_stop_step + 1, min_decode_steps) + if max_decode_steps < min_required_decode_steps: + raise RuntimeError( + "Invalid Ming decode window: " + f"max_decode_steps={max_decode_steps} is smaller than " + f"min_required_decode_steps={min_required_decode_steps}" + ) + should_force_stop = (step + 1) >= max_decode_steps + should_stop_head = ((step + 1) >= min_required_decode_steps) and stop_prob > stop_threshold + + if should_force_stop: + return ( + MING_STOP_REASON_MAX_DECODE_STEPS, + True, + True, + min_required_decode_steps, + text_eos_token_id, + ) + if should_stop_head: + return ( + MING_STOP_REASON_STOP_HEAD, + True, + False, + min_required_decode_steps, + text_eos_token_id, + ) + return ( + MING_STOP_REASON_CONTINUE, + False, + False, + min_required_decode_steps, + audio_dummy_token_id, + ) diff --git a/vllm_omni/model_executor/models/ming_tts/pipeline.py b/vllm_omni/model_executor/models/ming_tts/pipeline.py new file mode 100644 index 000000000..31ec98e5a --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/pipeline.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ming TTS pipeline: Stage-0 LLM+flow -> Stage-1 audio VAE.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.ming_tts" + +MING_TTS_PIPELINE = PipelineConfig( + model_type="ming_tts", + model_arch="MingTTSForConditionalGeneration", + hf_architectures=("MingTTSForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="llm", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + hf_config_name="llm_config", + engine_output_type="latent", + async_chunk_process_next_stage_input_func=(f"{_PROC}.llm2audio_vae_async_chunk"), + sampling_constraints={ + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 512, + "detokenize": True, + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="audio_vae", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + hf_config_name="llm_config", + engine_output_type="audio", + sync_process_input_func=f"{_PROC}.llm2audio_vae", + sampling_constraints={ + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 1, + "detokenize": False, + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py b/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py new file mode 100644 index 000000000..c732e07a0 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import copy +import json +import math +import re +from typing import Any + +import torch + +from .audio_prep import ( + coerce_speaker_embeddings, + count_prompt_latent_patches, + count_prompt_waveform_patches, + pad_prompt_waveform, +) +from .config_ming_tts import ( + AUDIO_FRAME_HOP, + KEY_MAX_DECODE_STEPS, + KEY_MIN_DECODE_STEPS, + KEY_PROMPT_LATENTS, + KEY_REQUEST_ID, + KEY_SPEAKER_EMBEDDING, + LATENT_DIM, + PATCH_SIZE, + SAMPLE_RATE, + VAE_PATCH_SIZE, +) + +DEFAULT_PROMPT = "Please generate speech based on the following description.\n" +BASE_CAPTION_TEMPLATE: dict[str, Any] = { + "audio_sequence": [ + { + "序号": 1, + "说话人": "speaker_1", + "方言": None, + "风格": None, + "语速": None, + "基频": None, + "音量": None, + "情感": None, + "BGM": { + "Genre": None, + "Mood": None, + "Instrument": None, + "Theme": None, + "ENV": None, + "SNR": None, + }, + "IP": None, + } + ] +} +_DURATION_SECONDS_RE = re.compile(r"Duration:\s*([0-9]+(?:\.[0-9]+)?)\s*s\b", re.IGNORECASE) + + +def create_instruction(user_input: Any) -> str | None: + if user_input is None: + return None + if isinstance(user_input, str): + return user_input + if not isinstance(user_input, dict): + raise ValueError(f"Ming instruction must be str, dict, or None; got {type(user_input).__name__}") + caption = copy.deepcopy(BASE_CAPTION_TEMPLATE) + item = caption["audio_sequence"][0] + for key, value in user_input.items(): + if key in item: + item[key] = value + return json.dumps(caption, ensure_ascii=False) + + +def parse_duration_seconds(text: str | None) -> float | None: + if not isinstance(text, str): + return None + match = _DURATION_SECONDS_RE.search(text) + if match is None: + return None + try: + value = float(match.group(1)) + except ValueError: + return None + if value <= 0.0: + return None + return value + + +def estimate_decode_steps_for_duration( + duration_seconds: float, + *, + sample_rate: int = SAMPLE_RATE, + frame_hop: int = AUDIO_FRAME_HOP, + patch_size: int = PATCH_SIZE, + vae_patch_size: int = VAE_PATCH_SIZE, +) -> int: + if duration_seconds <= 0.0: + return 0 + samples_per_decode_step = int(frame_hop) * int(patch_size) * int(vae_patch_size) + required_samples = float(duration_seconds) * float(sample_rate) + return max(1, int(math.ceil(required_samples / float(samples_per_decode_step)))) + + +def estimate_decode_step_window_for_duration(duration_seconds: float) -> tuple[int, int]: + target_steps = estimate_decode_steps_for_duration(duration_seconds) + min_steps = max(1, target_steps - 3) + max_steps = max(min_steps, target_steps + 3) + return min_steps, max_steps + + +def resolve_effective_runtime_controls( + *, + text: str, + runtime_controls: dict[str, Any] | None = None, +) -> dict[str, Any]: + controls = {} if runtime_controls is None else dict(runtime_controls) + has_explicit_min = KEY_MIN_DECODE_STEPS in controls and controls[KEY_MIN_DECODE_STEPS] is not None + has_explicit_max = KEY_MAX_DECODE_STEPS in controls and controls[KEY_MAX_DECODE_STEPS] is not None + if has_explicit_min or has_explicit_max: + return controls + duration_seconds = parse_duration_seconds(text) + if duration_seconds is None: + return controls + min_decode_steps, max_decode_steps = estimate_decode_step_window_for_duration(duration_seconds) + controls[KEY_MIN_DECODE_STEPS] = min_decode_steps + controls[KEY_MAX_DECODE_STEPS] = max_decode_steps + return controls + + +def build_dense_prompt_token_ids( + tokenizer: Any, + *, + prompt: str, + text: str, + instruction: str | None = None, + prompt_text: str | None = None, + speaker_count: int = 0, + prompt_patch_count: int = 0, +) -> list[int]: + speaker_prompt = [] + for idx in range(int(speaker_count)): + speaker_prompt.extend( + tokenizer.encode(f" speaker_{idx + 1}:") + + tokenizer.encode("<|vision_start|>") + + tokenizer.encode("<|vision_pad|>") + + tokenizer.encode("<|vision_end|>\n") + ) + instruction_prompt = ( + tokenizer.encode(instruction) + tokenizer.encode("<|endoftext|>") if instruction is not None else [] + ) + prompt_text_tokens = ( + tokenizer.encode(prompt_text) if int(prompt_patch_count) > 0 and prompt_text is not None else [] + ) + audio_patch_token_id = tokenizer.convert_tokens_to_ids("") + if audio_patch_token_id == tokenizer.unk_token_id: + raise ValueError("Ming tokenizer is missing required token.") + prompt_latent_tokens = [audio_patch_token_id] * int(prompt_patch_count) + text_input_prefix = ( + [] + if all(token in text for token in ("Genre: ", "Mood: ", "Instrument: ", "Theme: ", "Duration: ")) + else tokenizer.encode(" Text input:\n") + ) + return ( + tokenizer.encode("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n") + + tokenizer.encode("<|im_start|>user\n") + + tokenizer.encode(prompt) + + speaker_prompt + + text_input_prefix + + prompt_text_tokens + + tokenizer.encode(text) + + tokenizer.encode("<|im_end|>\n") + + tokenizer.encode("<|im_start|>assistant\n") + + instruction_prompt + + tokenizer.encode("